Update refacer.py

refacer.py

@@ -1,5 +1,3 @@
-import io
-import gradio as gr
 import cv2
 import onnxruntime as rt
 import sys
@@ -24,19 +22,20 @@ import re
 import subprocess
 
 class RefacerMode(Enum):
-    CPU, CUDA, COREML, TENSORRT = range(1, 5)
+     CPU, CUDA, COREML, TENSORRT = range(1, 5)
 
 class Refacer:
-    def __init__(self, force_cpu=False, colab_performance=False):
+    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:
+        if self.force_cpu :
             self.providers = ['CPUExecutionProvider']
         else:
             self.providers = rt.get_available_providers()
@@ -47,51 +46,62 @@ class Refacer:
 
         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)
+            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)
+            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)
+            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 the providers:
+            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))
+        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')
+        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 = 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)
+        self.face_swapper = INSwapper(model_path,sess_swap)
 
     def prepare_faces(self, faces):
-        self.replacement_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:
+                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:
@@ -99,13 +109,31 @@ class Refacer:
                 self.first_face = True
                 feat_original = None
                 print('No origin image: First face change')
-            _faces = self.__get_faces(face['destination'], max_num=1)
-            if len(_faces) < 1:
+            #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))
+            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):
 
-    def __get_faces(self, frame, max_num=0):
-        bboxes, kpss = self.face_detector.detect(frame, max_num=max_num, metric='default')
+        bboxes, kpss = self.face_detector.detect(frame,max_num=max_num,metric='default')
 
         if bboxes.shape[0] == 0:
             return []
@@ -121,91 +149,114 @@ class Refacer:
             ret.append(face)
         return ret
 
-    def process_first_face(self, frame):
-        faces = self.__get_faces(frame, max_num=1)
+    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)
+    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]:
+                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):
+    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):
-        results = []
-        with ThreadPoolExecutor(max_workers=self.use_num_cpus) as executor:
+        
+    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"))
+                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"))
-        return results
+                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))
 
-        frames = []
-        with tqdm(total=total_frames, desc="Extracting frames") as pbar:
+        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:
+                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()
 
-        refaced_frames = self.reface_group(faces, frames)
-
-        video_buffer = io.BytesIO()
-        out = cv2.VideoWriter('temp.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
-
-        for frame in refaced_frames:
-            out.write(frame)
-        out.release()
-
-        with open('temp.mp4', 'rb') as f:
-            video_buffer.write(f.read())
-        video_buffer.seek(0)
-
-        os.remove('temp.mp4')
-
-        return video_buffer
+        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'
 
-# Gradio Code
-def run(*vars):
-    video_path = vars[0]
-    origins = vars[1:(num_faces + 1)]
-    destinations = vars[(num_faces + 1):(num_faces * 2) + 1]
-    thresholds = vars[(num_faces * 2) + 1:]
+        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
 
-    faces = []
-    for k in range(0, num_faces):
-        if origins[k] is not None and destinations[k] is not None:
-            faces.append({
-                'origin': origins[k],
-                'destination[_{{{CITATION{{{_1{](https://github.com/qixinbo/OneButtonDeepLearning/tree/6e209f40102a7acaeb5d5798da013758c0ff9cd3/FaceSwap%2Fmenus%2FFaceSwap%2Finsightface_func%2Finsightface%2Fapp%2Fface_analysis.py)[_{{{CITATION{{{_2{](https://github.com/pgtinsley/arcface_aman/tree/7beda0d69dc40acc0138525ca84f50ecda126d8c/python-package%2Finsightface%2Fapp%2Fface_analysis.py)
+    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
+    }