Update refacer.py

refacer.py

@@ -1,3 +1,5 @@
+import io
+import gradio as gr
 import cv2
 import onnxruntime as rt
 import sys
@@ -103,21 +105,6 @@ class Refacer:
                 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"
-            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')
 
@@ -158,19 +145,18 @@ class Refacer:
         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):
+
+    def reface_group(self, faces, frames):
+        results = []
         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)
+        return results
 
     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)
@@ -181,48 +167,32 @@ class Refacer:
         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
+        frames = []
         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()
 
-        self.reface_group(faces,frames,output)
-        frames=[]
-        output.release()
+        refaced_frames = self.reface_group(faces, frames)
         
-        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'
-
-        if self.__try_ffmpeg_encoder('libx265'):
-            self.ffmpeg_video_encoder = 'libx265'
-        elif self.__try_ffmpeg_encoder('libvpx-vp9'):
-            self.ffmpeg_video_encoder = 'libvpx-vp9'
+        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