import cv2 import torch # hack to use onnxruntime with cuda import onnxruntime as rt import sys from insightface.app import FaceAnalysis sys.path.insert(1, './recognition') from scrfd import SCRFD from arcface_onnx import ArcFaceONNX import os from pathlib import Path from tqdm import tqdm import ffmpeg import random import multiprocessing as mp from concurrent.futures import ThreadPoolExecutor from insightface.model_zoo.inswapper import INSwapper import psutil from enum import Enum from insightface.app.common import Face from insightface.utils.storage import ensure_available import re import subprocess class RefacerMode(Enum): CPU, CUDA, COREML, TENSORRT = range(1, 5) class Refacer: def __init__(self,force_cpu=False,colab_performance=False): self.first_face = False self.force_cpu = force_cpu self.colab_performance = colab_performance self.__check_encoders() self.__check_providers() self.total_mem = psutil.virtual_memory().total self.__init_apps() def __check_providers(self): if self.force_cpu : self.providers = ['CPUExecutionProvider'] else: self.providers = rt.get_available_providers() rt.set_default_logger_severity(4) self.sess_options = rt.SessionOptions() self.sess_options.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL self.sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL if len(self.providers) == 1 and 'CPUExecutionProvider' in self.providers: self.mode = RefacerMode.CPU self.use_num_cpus = mp.cpu_count()-1 self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) print(f"CPU mode with providers {self.providers}") elif self.colab_performance: self.mode = RefacerMode.TENSORRT self.use_num_cpus = mp.cpu_count()-1 self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) print(f"TENSORRT mode with providers {self.providers}") elif 'CoreMLExecutionProvider' in self.providers: self.mode = RefacerMode.COREML self.use_num_cpus = mp.cpu_count()-1 self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) print(f"CoreML mode with providers {self.providers}") elif 'CUDAExecutionProvider' in self.providers: self.mode = RefacerMode.CUDA self.use_num_cpus = 2 self.sess_options.intra_op_num_threads = 1 if 'TensorrtExecutionProvider' in self.providers: self.providers.remove('TensorrtExecutionProvider') print(f"CUDA mode with providers {self.providers}") """ elif 'TensorrtExecutionProvider' in self.providers: self.mode = RefacerMode.TENSORRT #self.use_num_cpus = 1 #self.sess_options.intra_op_num_threads = 1 self.use_num_cpus = mp.cpu_count()-1 self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) print(f"TENSORRT mode with providers {self.providers}") """ def __init_apps(self): assets_dir = ensure_available('models', 'buffalo_l', root='~/.insightface') model_path = os.path.join(assets_dir, 'det_10g.onnx') sess_face = rt.InferenceSession(model_path, self.sess_options, providers=self.providers) self.face_detector = SCRFD(model_path,sess_face) self.face_detector.prepare(0,input_size=(640, 640)) model_path = os.path.join(assets_dir , 'w600k_r50.onnx') sess_rec = rt.InferenceSession(model_path, self.sess_options, providers=self.providers) self.rec_app = ArcFaceONNX(model_path,sess_rec) self.rec_app.prepare(0) model_path = 'inswapper_128.onnx' sess_swap = rt.InferenceSession(model_path, self.sess_options, providers=self.providers) self.face_swapper = INSwapper(model_path,sess_swap) def prepare_faces(self, faces): self.replacement_faces=[] for face in faces: #image1 = cv2.imread(face.origin) if "origin" in face: face_threshold = face['threshold'] bboxes1, kpss1 = self.face_detector.autodetect(face['origin'], max_num=1) if len(kpss1)<1: raise Exception('No face detected on "Face to replace" image') feat_original = self.rec_app.get(face['origin'], kpss1[0]) else: face_threshold = 0 self.first_face = True feat_original = None print('No origin image: First face change') #image2 = cv2.imread(face.destination) _faces = self.__get_faces(face['destination'],max_num=1) if len(_faces)<1: raise Exception('No face detected on "Destination face" image') self.replacement_faces.append((feat_original,_faces[0],face_threshold)) def __convert_video(self,video_path,output_video_path): if self.video_has_audio: print("Merging audio with the refaced video...") new_path = output_video_path + str(random.randint(0,999)) + "_c.mp4" #stream = ffmpeg.input(output_video_path) in1 = ffmpeg.input(output_video_path) in2 = ffmpeg.input(video_path) out = ffmpeg.output(in1.video, in2.audio, new_path,video_bitrate=self.ffmpeg_video_bitrate,vcodec=self.ffmpeg_video_encoder) out.run(overwrite_output=True,quiet=True) else: new_path = output_video_path print("The video doesn't have audio, so post-processing is not necessary") print(f"The process has finished.\nThe refaced video can be found at {os.path.abspath(new_path)}") return new_path def __get_faces(self,frame,max_num=0): bboxes, kpss = self.face_detector.detect(frame,max_num=max_num,metric='default') if bboxes.shape[0] == 0: return [] ret = [] for i in range(bboxes.shape[0]): bbox = bboxes[i, 0:4] det_score = bboxes[i, 4] kps = None if kpss is not None: kps = kpss[i] face = Face(bbox=bbox, kps=kps, det_score=det_score) face.embedding = self.rec_app.get(frame, kps) ret.append(face) return ret def process_first_face(self,frame): faces = self.__get_faces(frame,max_num=1) if len(faces) != 0: frame = self.face_swapper.get(frame, faces[0], self.replacement_faces[0][1], paste_back=True) return frame def process_faces(self,frame): faces = self.__get_faces(frame,max_num=0) for rep_face in self.replacement_faces: for i in range(len(faces) - 1, -1, -1): sim = self.rec_app.compute_sim(rep_face[0], faces[i].embedding) if sim>=rep_face[2]: frame = self.face_swapper.get(frame, faces[i], rep_face[1], paste_back=True) del faces[i] break return frame def __check_video_has_audio(self,video_path): self.video_has_audio = False probe = ffmpeg.probe(video_path) audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None) if audio_stream is not None: self.video_has_audio = True def reface_group(self, faces, frames, output): with ThreadPoolExecutor(max_workers = self.use_num_cpus) as executor: if self.first_face: results = list(tqdm(executor.map(self.process_first_face, frames), total=len(frames),desc="Processing frames")) else: results = list(tqdm(executor.map(self.process_faces, frames), total=len(frames),desc="Processing frames")) for result in results: output.write(result) def reface(self, video_path, faces): self.__check_video_has_audio(video_path) output_video_path = os.path.join('out',Path(video_path).name) self.prepare_faces(faces) cap = cv2.VideoCapture(video_path) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) print(f"Total frames: {total_frames}") fps = cap.get(cv2.CAP_PROP_FPS) frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fourcc = cv2.VideoWriter_fourcc(*'mp4v') output = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height)) frames=[] self.k = 1 with tqdm(total=total_frames,desc="Extracting frames") as pbar: while cap.isOpened(): flag, frame = cap.read() if flag and len(frame)>0: frames.append(frame.copy()) pbar.update() else: break if (len(frames) > 1000): self.reface_group(faces,frames,output) frames=[] cap.release() pbar.close() self.reface_group(faces,frames,output) frames=[] output.release() return self.__convert_video(video_path,output_video_path) def __try_ffmpeg_encoder(self, vcodec): print(f"Trying FFMPEG {vcodec} encoder") command = ['ffmpeg', '-y', '-f','lavfi','-i','testsrc=duration=1:size=1280x720:rate=30','-vcodec',vcodec,'testsrc.mp4'] try: subprocess.run(command, check=True, capture_output=True).stderr except subprocess.CalledProcessError as e: print(f"FFMPEG {vcodec} encoder doesn't work -> Disabled.") return False print(f"FFMPEG {vcodec} encoder works") return True def __check_encoders(self): self.ffmpeg_video_encoder='libx264' self.ffmpeg_video_bitrate='0' pattern = r"encoders: ([a-zA-Z0-9_]+(?: [a-zA-Z0-9_]+)*)" command = ['ffmpeg', '-codecs', '--list-encoders'] commandout = subprocess.run(command, check=True, capture_output=True).stdout result = commandout.decode('utf-8').split('\n') for r in result: if "264" in r: encoders = re.search(pattern, r).group(1).split(' ') for v_c in Refacer.VIDEO_CODECS: for v_k in encoders: if v_c == v_k: if self.__try_ffmpeg_encoder(v_k): self.ffmpeg_video_encoder=v_k self.ffmpeg_video_bitrate=Refacer.VIDEO_CODECS[v_k] print(f"Video codec for FFMPEG: {self.ffmpeg_video_encoder}") return VIDEO_CODECS = { 'h264_videotoolbox':'0', #osx HW acceleration 'h264_nvenc':'0', #NVIDIA HW acceleration #'h264_qsv', #Intel HW acceleration #'h264_vaapi', #Intel HW acceleration #'h264_omx', #HW acceleration 'libx264':'0' #No HW acceleration }