Spaces:
Running
Running
| import insightface | |
| import os | |
| import onnxruntime | |
| import cv2 | |
| import gfpgan | |
| import tempfile | |
| import time | |
| import gradio as gr | |
| class Predictor: | |
| def __init__(self): | |
| self.setup() | |
| def setup(self): | |
| os.makedirs('models', exist_ok=True) | |
| os.chdir('models') | |
| if not os.path.exists('GFPGANv1.4.pth'): | |
| os.system( | |
| 'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth' | |
| ) | |
| if not os.path.exists('inswapper_128.onnx'): | |
| os.system( | |
| 'wget https://huggingface.co/ashleykleynhans/inswapper/resolve/main/inswapper_128.onnx' | |
| ) | |
| os.chdir('..') | |
| """Load the model into memory to make running multiple predictions efficient""" | |
| self.face_swapper = insightface.model_zoo.get_model('models/inswapper_128.onnx', | |
| providers=onnxruntime.get_available_providers()) | |
| self.face_enhancer = gfpgan.GFPGANer(model_path='models/GFPGANv1.4.pth', upscale=1) | |
| self.face_analyser = insightface.app.FaceAnalysis(name='buffalo_l') | |
| self.face_analyser.prepare(ctx_id=0, det_size=(640, 640)) | |
| def get_face(self, img_data): | |
| analysed = self.face_analyser.get(img_data) | |
| try: | |
| largest = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) | |
| return largest | |
| except: | |
| print("No face found") | |
| return None | |
| def predict(self, input_image, swap_image): | |
| """Run a single prediction on the model""" | |
| try: | |
| frame = cv2.imread(input_image.name) | |
| face = self.get_face(frame) | |
| source_face = self.get_face(cv2.imread(swap_image.name)) | |
| try: | |
| print(frame.shape, face.shape, source_face.shape) | |
| except: | |
| print("printing shapes failed.") | |
| result = self.face_swapper.get(frame, face, source_face, paste_back=True) | |
| _, _, result = self.face_enhancer.enhance( | |
| result, | |
| paste_back=True | |
| ) | |
| out_path = tempfile.mkdtemp() + f"/{str(int(time.time()))}.jpg" | |
| cv2.imwrite(out_path, result) | |
| return out_path | |
| except Exception as e: | |
| print(f"{e}") | |
| return None | |
| # Instantiate the Predictor class | |
| predictor = Predictor() | |
| title = "Swap Faces Using Our Model!!!" | |
| # Create Gradio Interface | |
| iface = gr.Interface( | |
| fn=predictor.predict, | |
| inputs=[ | |
| gr.inputs.Image(type="file", label="Target Image"), | |
| gr.inputs.Image(type="file", label="Swap Image") | |
| ], | |
| outputs=gr.outputs.Image(type="file", label="Result"), | |
| title=title | |
| ) | |
| # Launch the Gradio Interface | |
| iface.launch() | |