multithreading

recognition/main.py

@@ -9,7 +9,7 @@ import onnxruntime
 from scrfd import SCRFD
 from arcface_onnx import ArcFaceONNX
 
-onnxruntime.set_default_logger_severity(3)
+onnxruntime.set_default_logger_severity(5)
 
 assets_dir = osp.expanduser('~/.insightface/models/buffalo_l')
 

refacer.py

@@ -12,11 +12,13 @@ from pathlib import Path
 from tqdm import tqdm
 import ffmpeg
 import random
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor
 
 class Refacer:
 
     def __init__(self):
-        onnxruntime.set_default_logger_severity(0)
+        onnxruntime.set_default_logger_severity(4)
 
         self.face_app = FaceAnalysis(name='buffalo_l')
         self.face_app.prepare(ctx_id=0, det_size=(640, 640))
@@ -37,9 +39,13 @@ class Refacer:
         for face in faces:
             #image1 = cv2.imread(face.origin)
             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])      
             #image2 = cv2.imread(face.destination)
             _faces = self.face_app.get(face['destination'],max_num=1)
+            if len(_faces)<1:
+                raise Exception('No face detected on "Destination face" image')
             replacements.append((feat_original,_faces[0],face['threshold']))
 
         return replacements
@@ -52,9 +58,18 @@ class Refacer:
         out.run()
         return new_path
     
+    def __process_faces(self,frame):
+        faces = self.face_app.get(frame)
+        for face in faces:
+            for rep_face in self.replacement_faces:
+                sim = self.rec_app.compute_sim(rep_face[0], face.embedding)
+                if sim>=rep_face[2]:
+                    frame = self.face_swapper.get(frame, face, rep_face[1], paste_back=True)
+        return frame
+
     def reface(self, video_path, faces):        
         output_video_path = os.path.join('out',Path(video_path).name)
-        replacement_faces=self.__prepare_faces(faces)
+        self.replacement_faces=self.__prepare_faces(faces)
 
         cap = cv2.VideoCapture(video_path)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -71,28 +86,23 @@ class Refacer:
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         output = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
         
-        with tqdm(total=total_frames) as pbar:
+        frames=[]
+        with tqdm(total=total_frames,desc="Extracting frames") as pbar:
             while cap.isOpened():
                 flag, frame = cap.read()
                 if flag and len(frame)>0:
-                    #pos_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+                    frames.append(frame.copy())
-                    faces = self.face_app.get(frame)
+                    pbar.update()
-                    res = frame.copy()
-
-                    for face in faces:
-                        for rep_face in replacement_faces:
-                            sim = self.rec_app.compute_sim(rep_face[0], face.embedding)
-                            if sim>=rep_face[2]:
-                                res = self.face_swapper.get(res, face, rep_face[1], paste_back=True)
-
-                    output.write(res)
-                    pbar.update(1)
                 else:
                     break
+            cap.release()
+            pbar.close()
         
-        pbar.close()
+        with ThreadPoolExecutor(max_workers = mp.cpu_count()-1) as executor:
-        cap.release()
+            results = list(tqdm(executor.map(self.__process_faces, frames), total=len(frames),desc="Processing frames"))
-        output.release()
+            for result in results:
+                output.write(result)
+            output.release()
 
         return self.__convert_video(video_path,output_video_path)