Rocky1/SadTalker

1import os2import cv23import yaml4import numpy as np5import warnings6from skimage import img_as_ubyte7 8warnings.filterwarnings('ignore')9 10import imageio11import torch12 13from src.facerender.modules.keypoint_detector import HEEstimator, KPDetector14from src.facerender.modules.mapping import MappingNet15from src.facerender.modules.generator import OcclusionAwareGenerator, OcclusionAwareSPADEGenerator16from src.facerender.modules.make_animation import make_animation 17 18from pydub import AudioSegment 19from src.utils.face_enhancer import enhancer as face_enhancer20from src.utils.paste_pic import paste_pic21from src.utils.videoio import save_video_with_watermark22 23 24class AnimateFromCoeff():25 26    def __init__(self, free_view_checkpoint, mapping_checkpoint,27                   config_path, device):28 29        with open(config_path) as f:30            config = yaml.safe_load(f)31 32        generator = OcclusionAwareSPADEGenerator(**config['model_params']['generator_params'],33                                                    **config['model_params']['common_params'])34        kp_extractor = KPDetector(**config['model_params']['kp_detector_params'],35                                    **config['model_params']['common_params'])36        he_estimator = HEEstimator(**config['model_params']['he_estimator_params'],37                               **config['model_params']['common_params'])38        mapping = MappingNet(**config['model_params']['mapping_params'])39 40 41        generator.to(device)42        kp_extractor.to(device)43        he_estimator.to(device)44        mapping.to(device)45        for param in generator.parameters():46            param.requires_grad = False47        for param in kp_extractor.parameters():48            param.requires_grad = False 49        for param in he_estimator.parameters():50            param.requires_grad = False51        for param in mapping.parameters():52            param.requires_grad = False53 54        if free_view_checkpoint is not None:55            self.load_cpk_facevid2vid(free_view_checkpoint, kp_detector=kp_extractor, generator=generator, he_estimator=he_estimator)56        else:57            raise AttributeError("Checkpoint should be specified for video head pose estimator.")58 59        if  mapping_checkpoint is not None:60            self.load_cpk_mapping(mapping_checkpoint, mapping=mapping)61        else:62            raise AttributeError("Checkpoint should be specified for video head pose estimator.") 63 64        self.kp_extractor = kp_extractor65        self.generator = generator66        self.he_estimator = he_estimator67        self.mapping = mapping68 69        self.kp_extractor.eval()70        self.generator.eval()71        self.he_estimator.eval()72        self.mapping.eval()73         74        self.device = device75    76    def load_cpk_facevid2vid(self, checkpoint_path, generator=None, discriminator=None, 77                        kp_detector=None, he_estimator=None, optimizer_generator=None, 78                        optimizer_discriminator=None, optimizer_kp_detector=None, 79                        optimizer_he_estimator=None, device="cpu"):80        checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))81        if generator is not None:82            generator.load_state_dict(checkpoint['generator'])83        if kp_detector is not None:84            kp_detector.load_state_dict(checkpoint['kp_detector'])85        if he_estimator is not None:86            he_estimator.load_state_dict(checkpoint['he_estimator'])87        if discriminator is not None:88            try:89               discriminator.load_state_dict(checkpoint['discriminator'])90            except:91               print ('No discriminator in the state-dict. Dicriminator will be randomly initialized')92        if optimizer_generator is not None:93            optimizer_generator.load_state_dict(checkpoint['optimizer_generator'])94        if optimizer_discriminator is not None:95            try:96                optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])97            except RuntimeError as e:98                print ('No discriminator optimizer in the state-dict. Optimizer will be not initialized')99        if optimizer_kp_detector is not None:100            optimizer_kp_detector.load_state_dict(checkpoint['optimizer_kp_detector'])101        if optimizer_he_estimator is not None:102            optimizer_he_estimator.load_state_dict(checkpoint['optimizer_he_estimator'])103 104        return checkpoint['epoch']105    106    def load_cpk_mapping(self, checkpoint_path, mapping=None, discriminator=None,107                 optimizer_mapping=None, optimizer_discriminator=None, device='cpu'):108        checkpoint = torch.load(checkpoint_path,  map_location=torch.device(device))109        if mapping is not None:110            mapping.load_state_dict(checkpoint['mapping'])111        if discriminator is not None:112            discriminator.load_state_dict(checkpoint['discriminator'])113        if optimizer_mapping is not None:114            optimizer_mapping.load_state_dict(checkpoint['optimizer_mapping'])115        if optimizer_discriminator is not None:116            optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])117 118        return checkpoint['epoch']119 120    def generate(self, x, video_save_dir, pic_path, crop_info, enhancer=None, background_enhancer=None, preprocess='crop'):121 122        source_image=x['source_image'].type(torch.FloatTensor)123        source_semantics=x['source_semantics'].type(torch.FloatTensor)124        target_semantics=x['target_semantics_list'].type(torch.FloatTensor) 125        source_image=source_image.to(self.device)126        source_semantics=source_semantics.to(self.device)127        target_semantics=target_semantics.to(self.device)128        if 'yaw_c_seq' in x:129            yaw_c_seq = x['yaw_c_seq'].type(torch.FloatTensor)130            yaw_c_seq = x['yaw_c_seq'].to(self.device)131        else:132            yaw_c_seq = None133        if 'pitch_c_seq' in x:134            pitch_c_seq = x['pitch_c_seq'].type(torch.FloatTensor)135            pitch_c_seq = x['pitch_c_seq'].to(self.device)136        else:137            pitch_c_seq = None138        if 'roll_c_seq' in x:139            roll_c_seq = x['roll_c_seq'].type(torch.FloatTensor) 140            roll_c_seq = x['roll_c_seq'].to(self.device)141        else:142            roll_c_seq = None143 144        frame_num = x['frame_num']145 146        predictions_video = make_animation(source_image, source_semantics, target_semantics,147                                        self.generator, self.kp_extractor, self.he_estimator, self.mapping, 148                                        yaw_c_seq, pitch_c_seq, roll_c_seq, use_exp = True)149 150        predictions_video = predictions_video.reshape((-1,)+predictions_video.shape[2:])151        predictions_video = predictions_video[:frame_num]152 153        video = []154        for idx in range(predictions_video.shape[0]):155            image = predictions_video[idx]156            image = np.transpose(image.data.cpu().numpy(), [1, 2, 0]).astype(np.float32)157            video.append(image)158        result = img_as_ubyte(video)159 160        ### the generated video is 256x256, so we  keep the aspect ratio, 161        original_size = crop_info[0]162        if original_size:163            result = [ cv2.resize(result_i,(256, int(256.0 * original_size[1]/original_size[0]) )) for result_i in result ]164        165        video_name = x['video_name']  + '.mp4'166        path = os.path.join(video_save_dir, 'temp_'+video_name)167        imageio.mimsave(path, result, fps=float(25))168 169        av_path = os.path.join(video_save_dir, video_name)170        return_path = av_path 171        172        audio_path =  x['audio_path'] 173        audio_name = os.path.splitext(os.path.split(audio_path)[-1])[0]174        new_audio_path = os.path.join(video_save_dir, audio_name+'.wav')175        start_time = 0176        sound = AudioSegment.from_mp3(audio_path)177        frames = frame_num 178        end_time = start_time + frames*1/25*1000179        word1=sound.set_frame_rate(16000)180        word = word1[start_time:end_time]181        word.export(new_audio_path, format="wav")182 183        save_video_with_watermark(path, new_audio_path, av_path, watermark= None)184        print(f'The generated video is named {video_name} in {video_save_dir}')185 186        if preprocess.lower() == 'full':187            # only add watermark to the full image.188            video_name_full = x['video_name']  + '_full.mp4'189            full_video_path = os.path.join(video_save_dir, video_name_full)190            return_path = full_video_path191            paste_pic(path, pic_path, crop_info, new_audio_path, full_video_path)192            print(f'The generated video is named {video_save_dir}/{video_name_full}') 193        else:194            full_video_path = av_path 195 196        #### paste back then enhancers197        if enhancer:198            video_name_enhancer = x['video_name']  + '_enhanced.mp4'199            enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer)200            av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) 201            return_path = av_path_enhancer202            enhanced_images = face_enhancer(full_video_path, method=enhancer, bg_upsampler=background_enhancer)203            imageio.mimsave(enhanced_path, enhanced_images, fps=float(25))204            205            save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= None)206            print(f'The generated video is named {video_save_dir}/{video_name_enhancer}')207            os.remove(enhanced_path)208 209        os.remove(path)210        os.remove(new_audio_path)211 212        return return_path213 214