undercovercd/Swap-Face-Model

1import insightface2import os3import onnxruntime4import cv25import gfpgan6import tempfile7import time8import gradio as gr9 10 11class Predictor:12    def __init__(self):13        self.setup()14 15    def setup(self):16        os.makedirs('models', exist_ok=True)17        os.chdir('models')18        if not os.path.exists('GFPGANv1.4.pth'):19            os.system(20                'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth'21            )22        if not os.path.exists('inswapper_128.onnx'):23            os.system(24                'wget https://huggingface.co/ashleykleynhans/inswapper/resolve/main/inswapper_128.onnx'25            )26        os.chdir('..')27 28        """Load the model into memory to make running multiple predictions efficient"""29        self.face_swapper = insightface.model_zoo.get_model('models/inswapper_128.onnx',30                                                            providers=onnxruntime.get_available_providers())31        self.face_enhancer = gfpgan.GFPGANer(model_path='models/GFPGANv1.4.pth', upscale=1)32        self.face_analyser = insightface.app.FaceAnalysis(name='buffalo_l')33        self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))34 35    def get_face(self, img_data):36        analysed = self.face_analyser.get(img_data)37        try:38            largest = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))39            return largest40        except:41            print("No face found")42            return None43 44    def predict(self, input_image, swap_image):45        """Run a single prediction on the model"""46        try:47            frame = cv2.imread(input_image.name)48            face = self.get_face(frame)49            source_face = self.get_face(cv2.imread(swap_image.name))50            try:51                print(frame.shape, face.shape, source_face.shape)52            except:53                print("printing shapes failed.")54            result = self.face_swapper.get(frame, face, source_face, paste_back=True)55 56            _, _, result = self.face_enhancer.enhance(57                result,58                paste_back=True59            )60            out_path = tempfile.mkdtemp() + f"/{str(int(time.time()))}.jpg"61            cv2.imwrite(out_path, result)62            return out_path63        except Exception as e:64            print(f"{e}")65            return None66 67 68# Instantiate the Predictor class69predictor = Predictor()70title = "Swap Faces Using Our Model!!!"71 72# Create Gradio Interface73iface = gr.Interface(74    fn=predictor.predict,75    inputs=[76        gr.inputs.Image(type="file", label="Target Image"),77        gr.inputs.Image(type="file", label="Swap Image")78    ],79    outputs=gr.outputs.Image(type="file", label="Result"),80    title=title,81    examples=[["input.jpg", "swap img.jpg"]])82 83 84# Launch the Gradio Interface85iface.launch()86