| | 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('..') |
| |
|
| | 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): |
| | try: |
| | frame = cv2.imread(input_image.name) |
| | face = self.get_face(frame) |
| | source_face = self.get_face(cv2.imread(swap_image.name)) |
| | 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() |
| |
|
| | css = """ |
| | footer { |
| | visibility: hidden; |
| | } |
| | """ |
| |
|
| | |
| | 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"), |
| | examples=[["input.jpg", "swap img.jpg"]], |
| | css=css |
| | ) |
| |
|
| | |
| | iface.launch() |
| |
|
