Update refacer.py

refacer.py

@@ -1,49 +1,87 @@
 import cv2
-import ffmpeg
-import numpy as np
-import os
 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.path as osp
+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):
+    def __init__(self,force_cpu=False,colab_performance=False):
+        self.first_face = False
         self.force_cpu = force_cpu
         self.colab_performance = colab_performance
-        self.sess_options = None
-        self.providers = ['CPUExecutionProvider']
-        self.model_path = "./inswapper_128.onnx"  # Ensure the model path is correct
+        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}")
+
     def __init_apps(self):
-        self.sess = rt.InferenceSession(self.model_path, self.sess_options, providers=self.providers)
-
-    def __convert_video(self, video_path, output_video_path):
-        """
-        Convert the video using ffmpeg and return the output file.
-        """
-        try:
-            # Use ffmpeg to convert the video
-            out = ffmpeg.input(video_path)
-            out = ffmpeg.output(out, output_video_path, vcodec='libx264', acodec='aac')  # Ensure appropriate codecs
-            out.run(overwrite_output=True, quiet=False)  # verbose output to capture any errors
-        except ffmpeg.Error as e:
-            print("FFmpeg Error:", e.stderr.decode())  # Print detailed error information
-            raise
-        return output_video_path
-
-    def reface(self, video_path, faces):
-        """
-        Refacing the video by swapping faces.
-        """
-        output_video_path = "output_video.mp4"  # Set the output video path
-        # Code to process the faces and perform refacing (image manipulation here)
-        # Process faces and apply face swap logic
-        print("Refacing video...")
-        return self.__convert_video(video_path, output_video_path)
-
-# Example of running the app (just for understanding, Gradio app will call this part)
-if __name__ == "__main__":
-    refacer = Refacer(force_cpu=True)
-    faces = [{'origin': 'origin_image.jpg', 'destination': 'destination_image.jpg', 'threshold': 0.8}]
-    video_path = 'input_video.mp4'
-    refaced_video_path = refacer.reface(video_path, faces)
-    print(f"Refaced video can be found at: {refaced_video_path}")
+        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)
+
+    # மேலும் தேவையான முறைமைகள்...