MakeItTalk/src/dataset/utils/Av2Flau_Convertor.py

"""
 # Copyright 2020 Adobe
 # All Rights Reserved.
 
 # NOTICE: Adobe permits you to use, modify, and distribute this file in
 # accordance with the terms of the Adobe license agreement accompanying
 # it.
 
"""

import numpy as np
import os
import ffmpeg
import cv2
import face_alignment
from src.dataset.utils import icp


class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y


class ShapeParts:
    def __init__(self, np_pts):
        self.data = np_pts

    def part(self, idx):
        return Point(self.data[idx, 0], self.data[idx, 1])


class Av2Flau_Convertor():
    """

    Any video to facial landmark and audio numpy data converter.

    """

    def __init__(self, video_dir, out_dir, idx=0):

        self.video_dir = video_dir
        if ('\\' in video_dir):
            self.video_name = video_dir.split('\\')[-1]
        else:
            self.video_name = video_dir.split('/')[-1]
        self.out_dir = out_dir
        self.idx = idx
        self.input_format = self.video_dir[-4:]

        # landmark predictor = FANet
        self.predictor = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device='cuda', flip_input=True)

        # landmark register
        self.t_shape_idx = (27, 28, 29, 30, 33, 36, 39, 42, 45)

    def convert(self, max_num_frames=250, save_audio=False, show=False, register=False):

        # Step 1: preclean video: check stream==2, convert fps/sample_rate,
        ret, wfn = self.__preclean_video__()
        if (not ret):
            return

        # Step 2: detect facial landmark
        wfn = self.video_dir.replace(self.input_format, '_preclean.mp4')
        ret, fl2d, fl3d = self.__video_facial_landmark_detection__(video_dir=wfn, display=False, max_num_frames=max_num_frames)
        if (not ret):
            return
        if (len(fl3d) < 9):
            print('The length of the landmark is too short, skip')
            return

        # Step 3: raw save landmark / audio
        fl3d = np.array(fl3d)
        np.savetxt(os.path.join(self.out_dir, 'raw_fl3d/fan_{:05d}_{}_3d.txt'.format(self.idx, self.video_name[:-4])),
                   fl3d, fmt='%.2f')
        if (save_audio):
            self.__save_audio__(video_dir=self.video_dir.replace(self.input_format, '_preclean.mp4'), fl3d=fl3d)

        # Step 3.5: merge a/v together (optional)
        if (show):
            sf, ef = (fl3d[0][0], fl3d[-1][0]) if fl3d.shape[0] > 0 else (0, 0)
            print(sf, ef)
            print(self.video_dir.replace(self.input_format, '_fl_detect.mp4'),
                  os.path.join(self.out_dir, 'tmp_v', '{:05d}_{}_fl_av.mp4'.format(
                      self.idx, self.video_name[:-4]))
                  )
            self.__ffmpeg_merge_av__(
                video_dir=self.video_dir.replace(self.input_format, '_fl_detect.mp4'),
                audio_dir=self.video_dir.replace(self.input_format, '_preclean.mp4'),
                WriteFileName=os.path.join(self.out_dir, 'tmp_v', '{:05d}_{}_fl_av.mp4'.format(
                    self.idx, self.video_name[:-4])),
                start_end_frame=(int(sf), int(ef)))

        # Step 4: remove tmp files
        os.remove(self.video_dir.replace(self.input_format, '_preclean.mp4'))
        if(os.path.isfile(self.video_dir.replace(self.input_format, '_fl_detect.mp4'))):
            os.remove(self.video_dir.replace(self.input_format, '_fl_detect.mp4'))

        # Step 5: register fl3d
        if (register):
            self.__single_landmark_3d_register__(fl3d)
            # TODO: visualize register fl3d

    ''' ========================================================================

                            STEP 1: Preclean video

    ======================================================================== '''

    def __preclean_video__(self, WriteFileName='_preclean.mp4', fps=25, sample_rate=16000):
        '''
        Pre-clean downloaded videos. Return false if more than 2 streams found.
        Then convert it to fps=25, sample_rate=16kHz
        '''
        input_video_dir = self.video_dir if '_x_' not in self.video_dir else self.video_dir.replace('_x_', '/')

        probe = ffmpeg.probe(input_video_dir)
        # print(probe['streams'])
        # print(len(probe['streams']))
        # if(len(probe['streams']) != 2):
        #     print('Error: not valid for # of a/v channel == 2.')
        #     return False, None
        # exit(0)
        # probe['streams'] = probe['streams'][0::2]

        codec = {'video': '', 'audio': ''}
        for i, stream in enumerate(probe['streams'][0:2]):
            codec[stream['codec_type']] = stream['codec_name']

        # create preclean video
        (
            ffmpeg
                .input(input_video_dir)
                .output(self.video_dir.replace(self.input_format, WriteFileName),
                        # vcodec=codec['video'],
                        # acodec=codec['audio'],
                        r=fps, ar=sample_rate)
                .overwrite_output().global_args('-loglevel', 'quiet')
                .run()
        )

        return True, self.video_dir.replace(self.input_format, WriteFileName)

    ''' ========================================================================

                       STEP 2: Detect facial landmark

    ======================================================================== '''

    def __video_facial_landmark_detection__(self, video_dir=None, display=False, WriteFileName='_fl_detect.mp4',
                                            max_num_frames=250, write=False):
        '''
        Get facial landmark from video.
        '''

        # load video
        print('video_dir : ' + video_dir)
        video = cv2.VideoCapture(video_dir)

        # return false if cannot open
        if (video.isOpened() == False):
            print('Unable to open video file')
            return False, None

        # display info
        length = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = video.get(cv2.CAP_PROP_FPS)
        w = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print('Process Video {}, len: {}, FPS: {:.2f}, W X H: {} x {}'.format(video_dir, length, fps, w, h))

        if(write):
            writer = cv2.VideoWriter(self.video_dir.replace(self.input_format, WriteFileName),
                                 cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (w, h))

        video_facial_landmark = []  # face-landmark np array per frame =: idx + [x,y] * 68
        video_facial_landmark_3d = []  # face-landmark np array per frame =: idx + [x,y,z] * 68
        frame_id = 0
        not_detected_frames = 0

        while (video.isOpened()):
            ret, frame = video.read()
            # reach EOF
            if (ret == False):
                break

            # too many not-detected frames (in middle of video)
            if (not_detected_frames > 5):
                if (len(video_facial_landmark) < 10):
                    # at beginning of the video
                    video_facial_landmark = []
                    video_facial_landmark_3d = []
                else:
                    break

            # dlib facial landmark detect
            img_ret, shape, shape_3d = self.__image_facial_landmark_detection__(img=frame)

            # successfully detected
            if (img_ret):
                # print('\t ==> frame {}/{}'.format(frame_id, length))

                # current frame xy coordinates
                xys = []
                for part_i in range(68):
                    xys.append(shape.part(part_i).x)
                    xys.append(shape.part(part_i).y)

                # check any not_detected_frames, and interp them
                if (not_detected_frames > 0 and len(video_facial_landmark) > 0):
                    # interpolate
                    def interp(last, cur, num, dims=68 * 2 + 1):
                        interp_xys_np = np.zeros((num, dims))
                        for dim in range(dims):
                            interp_xys_np[:, dim] = np.interp(np.arange(0, num), [-1, num], [last[dim], cur[dim]])
                        interp_xys_np = np.round(interp_xys_np).astype('int')
                        interp_xys = [list(xy) for xy in interp_xys_np]
                        return interp_xys

                    interp_xys = interp(video_facial_landmark[-1], [frame_id] + xys, not_detected_frames)
                    video_facial_landmark += interp_xys

                not_detected_frames = 0

                # save landmark/frame_index
                video_facial_landmark.append([frame_id] + xys)
                if (shape_3d.any()):
                    video_facial_landmark_3d.append([frame_id] + list(np.reshape(shape_3d, -1)))

                if(write):
                    frame = self.__vis_landmark_on_img__(frame, shape)

            else:
                print('\t ==> frame {}/{} Not detected'.format(frame_id, length))
                not_detected_frames += 1

            if (display):
                cv2.imshow('Frame', frame)
                if (cv2.waitKey(10) == ord('q')):
                    break

            if(write):
                writer.write(frame)
            frame_id += 1

            if(frame_id > max_num_frames):
                break

        video.release()
        if(write):
            writer.release()
        cv2.destroyAllWindows()

        print('\t ==> Final processed frames {}/{}'.format(frame_id, length))

        return True, video_facial_landmark, video_facial_landmark_3d

    def __image_facial_landmark_detection__(self, img=None):
        '''
        Get facial landmark from single image by FANet
        '''

        shapes = self.predictor.get_landmarks(img)
        if (not shapes):
            return False, None, None

        max_size_idx = 0
        shape = ShapeParts(shapes[max_size_idx][:, 0:2])
        shape_3d = shapes[max_size_idx]

        # when use 2d estimator
        shape_3d = np.concatenate([shape_3d, np.ones(shape=(68, 1))], axis=1)

        return True, shape, shape_3d

    def __vis_landmark_on_img__(self, img, shape, linewidth=2):
        '''
        Visualize landmark on images.
        '''
        if (type(shape) == ShapeParts):
            def draw_curve(idx_list, color=(0, 255, 0), loop=False, lineWidth=linewidth):
                for i in idx_list:
                    cv2.line(img, (shape.part(i).x, shape.part(i).y), (shape.part(i + 1).x, shape.part(i + 1).y),
                             color, lineWidth)
                if (loop):
                    cv2.line(img, (shape.part(idx_list[0]).x, shape.part(idx_list[0]).y),
                             (shape.part(idx_list[-1] + 1).x, shape.part(idx_list[-1] + 1).y), color, lineWidth)

            draw_curve(list(range(0, 16)))  # jaw
            draw_curve(list(range(17, 21)))  # eye brow
            draw_curve(list(range(22, 26)))
            draw_curve(list(range(27, 35)))  # nose
            draw_curve(list(range(36, 41)), loop=True)  # eyes
            draw_curve(list(range(42, 47)), loop=True)
            draw_curve(list(range(48, 59)), loop=True)  # mouth
            draw_curve(list(range(60, 67)), loop=True)

        else:
            def draw_curve(idx_list, color=(0, 255, 0), loop=False, lineWidth=linewidth):
                for i in idx_list:
                    cv2.line(img, (shape[i, 0], shape[i, 1]), (shape[i + 1, 0], shape[i + 1, 1]), color, lineWidth)
                if (loop):
                    cv2.line(img, (shape[idx_list[0], 0], shape[idx_list[0], 1]),
                             (shape[idx_list[-1] + 1, 0], shape[idx_list[-1] + 1, 1]), color, lineWidth)

            draw_curve(list(range(0, 16)))  # jaw
            draw_curve(list(range(17, 21)))  # eye brow
            draw_curve(list(range(22, 26)))
            draw_curve(list(range(27, 35)))  # nose
            draw_curve(list(range(36, 41)), loop=True)  # eyes
            draw_curve(list(range(42, 47)), loop=True)
            draw_curve(list(range(48, 59)), loop=True)  # mouth
            draw_curve(list(range(60, 67)), loop=True)

        return img

    def __ffmpeg_merge_av__(self, video_dir, audio_dir, WriteFileName, start_end_frame):
        probe = ffmpeg.probe(video_dir)
        fps = probe['streams'][0]['avg_frame_rate']
        spf = float(fps.split('/')[1]) / float(fps.split('/')[0])
        sf, ef = start_end_frame
        st, tt = sf * spf, ef * spf - sf * spf

        vin = ffmpeg.input(video_dir).video
        # ain = ffmpeg.input(audio_dir).audio
        # out = ffmpeg.output(vin, ain, WriteFileName, codec='copy', ss=st, t=tt, shortest=None)
        out = ffmpeg.output(vin, WriteFileName, codec='copy', ss=st, t=tt, shortest=None)
        out = out.overwrite_output().global_args('-loglevel', 'quiet')
        out.run()

        # os.system('ffmpeg -i {} -codec copy -ss {} -t {} {}'.format(video_dir, st, tt, WriteFileName))

    def __save_audio__(self, video_dir, fl3d):
        """
        Extract audio from preclean video. Used for creating audio-aware dataset.

        """
        sf, ef = fl3d[0][0], fl3d[-1][0]

        probe = ffmpeg.probe(video_dir)
        fps = probe['streams'][0]['avg_frame_rate']
        spf = float(fps.split('/')[1]) / float(fps.split('/')[0])
        st, tt = sf * spf, ef * spf - sf * spf

        audio_dir = os.path.join(self.out_dir, 'raw_wav', '{:05d}_{}_audio.wav'.format(self.idx, self.video_name[:-4]))
        (
            ffmpeg
                .input(video_dir)
                .output(audio_dir, ss=st, t=tt)
                .overwrite_output().global_args('-loglevel', 'quiet')
                .run()
        )

    ''' ========================================================================

                            STEP 5: Landmark register

    ======================================================================== '''

    def __single_landmark_3d_register__(self, fl3d, display=False):
        """
        Register a single 3d landmark file

        """
        # Step 1 : Load and Smooth
        from scipy.signal import savgol_filter
        lines = savgol_filter(fl3d, 7, 3, axis=0)

        all_landmarks = lines[:, 1:].reshape((-1, 68, 3))  # remove frame idx
        w, h = int(np.max(all_landmarks[:, :, 0])) + 20, int(np.max(all_landmarks[:, :, 1])) + 20

        # Step 2 : setup anchor face
        print('Using exisiting ' + 'dataset/utils/ANCHOR_T_SHAPE_{}.txt'.format(len(self.t_shape_idx)))
        anchor_t_shape = np.loadtxt('dataset/utils/ANCHOR_T_SHAPE_{}.txt'.format(len(self.t_shape_idx)))

        registered_landmarks_to_save = []
        registered_affine_mat_to_save = []
        # for each line
        for line in lines:
            frame_id = line[0]
            landmarks = line[1:].reshape(68, 3)

            # Step 3 : ICP on (frame, anchor)
            frame_t_shape = landmarks[self.t_shape_idx, :]

            T, distance, itr = icp(frame_t_shape, anchor_t_shape)

            # Step 4 : Affine transform
            landmarks = np.hstack((landmarks, np.ones((68, 1))))
            registered_landmarks = np.dot(T, landmarks.T).T
            err = np.mean(np.sqrt(np.sum((registered_landmarks[self.t_shape_idx, 0:3] - anchor_t_shape) ** 2, axis=1)))
            # print(err, distance, itr)

            # Step 5 : Save is requested
            registered_landmarks_to_save.append([frame_id] + list(registered_landmarks[:, 0:3].reshape(-1)))
            registered_affine_mat_to_save.append([frame_id] + list(T.reshape(-1)))

            # Step 5.5 (optional) : visualize ori / registered faces (Isolated in Black BG)
            if (display):
                img = np.zeros((h, w * 2, 3), np.uint8)
                self.__vis_landmark_on_img__(img, landmarks.astype(np.int))
                registered_landmarks[:, 0] += w
                self.__vis_landmark_on_img__(img, registered_landmarks.astype(np.int))
                cv2.imshow('img', img)
                if (cv2.waitKey(30) == ord('q')):
                    break

        np.savetxt(os.path.join(self.out_dir, 'register_fl3d', '{:05d}_{}_fl_sm.txt'
                                .format(self.idx, self.video_name[:-4])),
                   lines, fmt='%.6f')
        np.savetxt(os.path.join(self.out_dir, 'register_fl3d', '{:05d}_{}_fl_reg.txt'
                                .format(self.idx, self.video_name[:-4])),
                   np.array(registered_landmarks_to_save), fmt='%.6f')
        np.savetxt(os.path.join(self.out_dir, 'register_fl3d', '{:05d}_{}_mat_reg.txt'
                                .format(self.idx, self.video_name[:-4])),
                   np.array(registered_affine_mat_to_save), fmt='%.6f')


if __name__ == '__main__':
    video_dir = r'C:\Users\yangzhou\Videos\004_1.mp4'
    out_dir = r'C:\Users\yangzhou\Videos'
    c = Av2Flau_Convertor(video_dir, out_dir, idx=0)
    c.convert()