import sys import argparse import cv2 import torch import time import os from utils.inference.image_processing import crop_face, get_final_image from utils.inference.video_processing import read_video, get_target, get_final_video, add_audio_from_another_video, face_enhancement from utils.inference.core import model_inference from network.AEI_Net import AEI_Net from coordinate_reg.image_infer import Handler from insightface_func.face_detect_crop_multi import Face_detect_crop from arcface_model.iresnet import iresnet100 from models.pix2pix_model import Pix2PixModel from models.config_sr import TestOptions def init_models(args): # model for face cropping app = Face_detect_crop(name='antelope', root='./insightface_func/models') app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640)) # main model for generation G = AEI_Net(args.backbone, num_blocks=args.num_blocks, c_id=512) G.eval() G.load_state_dict(torch.load(args.G_path, map_location=torch.device('cpu'))) G = G.cuda() G = G.half() # arcface model to get face embedding netArc = iresnet100(fp16=False) netArc.load_state_dict(torch.load('arcface_model/backbone.pth')) netArc=netArc.cuda() netArc.eval() # model to get face landmarks handler = Handler('./coordinate_reg/model/2d106det', 0, ctx_id=0, det_size=640) # model to make superres of face, set use_sr=True if you want to use super resolution or use_sr=False if you don't if args.use_sr: os.environ['CUDA_VISIBLE_DEVICES'] = '0' torch.backends.cudnn.benchmark = True opt = TestOptions() #opt.which_epoch ='10_7' model = Pix2PixModel(opt) model.netG.train() else: model = None return app, G, netArc, handler, model def main(args): app, G, netArc, handler, model = init_models(args) # get crops from source images print('List of source paths: ',args.source_paths) source = [] try: for source_path in args.source_paths: img = cv2.imread(source_path) img = crop_face(img, app, args.crop_size)[0] source.append(img[:, :, ::-1]) except TypeError: print("Bad source images!") exit() # get full frames from video if not args.image_to_image: full_frames, fps = read_video(args.target_video) else: target_full = cv2.imread(args.target_image) full_frames = [target_full] # get target faces that are used for swap set_target = True print('List of target paths: ', args.target_faces_paths) if not args.target_faces_paths: target = get_target(full_frames, app, args.crop_size) set_target = False else: target = [] try: for target_faces_path in args.target_faces_paths: img = cv2.imread(target_faces_path) img = crop_face(img, app, args.crop_size)[0] target.append(img) except TypeError: print("Bad target images!") exit() start = time.time() final_frames_list, crop_frames_list, full_frames, tfm_array_list = model_inference(full_frames, source, target, netArc, G, app, set_target, similarity_th=args.similarity_th, crop_size=args.crop_size, BS=args.batch_size) if args.use_sr: final_frames_list = face_enhancement(final_frames_list, model) if not args.image_to_image: get_final_video(final_frames_list, crop_frames_list, full_frames, tfm_array_list, args.out_video_name, fps, handler) add_audio_from_another_video(args.target_video, args.out_video_name, "audio") print(f"Video saved with path {args.out_video_name}") else: result = get_final_image(final_frames_list, crop_frames_list, full_frames[0], tfm_array_list, handler) cv2.imwrite(args.out_image_name, result) print(f'Swapped Image saved with path {args.out_image_name}') print('Total time: ', time.time()-start) if __name__ == "__main__": parser = argparse.ArgumentParser() # Generator params parser.add_argument('--G_path', default='weights/G_unet_2blocks.pth', type=str, help='Path to weights for G') parser.add_argument('--backbone', default='unet', const='unet', nargs='?', choices=['unet', 'linknet', 'resnet'], help='Backbone for attribute encoder') parser.add_argument('--num_blocks', default=2, type=int, help='Numbers of AddBlocks at AddResblock') parser.add_argument('--batch_size', default=40, type=int) parser.add_argument('--crop_size', default=224, type=int, help="Don't change this") parser.add_argument('--use_sr', default=False, type=bool, help='True for super resolution on swap images') parser.add_argument('--similarity_th', default=0.15, type=float, help='Threshold for selecting a face similar to the target') parser.add_argument('--source_paths', default=['examples/images/mark.jpg', 'examples/images/elon_musk.jpg'], nargs='+') parser.add_argument('--target_faces_paths', default=[], nargs='+', help="It's necessary to set the face/faces in the video to which the source face/faces is swapped. You can skip this parametr, and then any face is selected in the target video for swap.") # parameters for image to video parser.add_argument('--target_video', default='examples/videos/nggyup.mp4', type=str, help="It's necessary for image to video swap") parser.add_argument('--out_video_name', default='examples/results/result.mp4', type=str, help="It's necessary for image to video swap") # parameters for image to image parser.add_argument('--image_to_image', default=False, type=bool, help='True for image to image swap, False for swap on video') parser.add_argument('--target_image', default='examples/images/beckham.jpg', type=str, help="It's necessary for image to image swap") parser.add_argument('--out_image_name', default='examples/results/result.png', type=str,help="It's necessary for image to image swap") args = parser.parse_args() main(args)