| 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 |
|
|
|
|
| |
| predictor = Predictor() |
| title = "Swap Faces Using Our Model!!!" |
|
|
| |
| 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, |
| examples=[["input.jpg", "swap img.jpg"]]) |
|
|
|
|
| |
| iface.launch() |
|
|
