Update refacer.py

refacer.py

@@ -1,262 +1,44 @@
+import os
 import cv2
-import onnxruntime as rt
-import sys
+import numpy as np
 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)
+from insightface.model_zoo import model_zoo
+from onnxruntime import InferenceSession
 
 class Refacer:
-    def __init__(self,force_cpu=False,colab_performance=False):
-        self.first_face = False
+    def __init__(self, force_cpu=False, colab_performance=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)
+        self.model = self.load_model()
 
-    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 load_model(self):
+        model_path = "/home/user/app/inswapper_128.onnx"  # Replace with your actual model path
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"Model not found at {model_path}")
+        return InferenceSession(model_path, providers=["CPUExecutionProvider" if self.force_cpu else "CUDAExecutionProvider"])
 
     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)
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found at {video_path}")
 
         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=[]
+        output_path = "output_video.mp4"
+        out = cv2.VideoWriter(output_path, fourcc, cap.get(cv2.CAP_PROP_FPS),
+                              (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
 
-            cap.release()
-            pbar.close()
+            for face in faces:
+                # Here, you should apply face-swapping logic using ONNX and the destination face
+                pass  # Replace this with face-swapping code
 
-        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'
+            out.write(frame)
 
-        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
+        cap.release()
+        out.release()
 
-    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
-    }
+        return output_path