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

    # 这是修正后的predict函数，请用它替换掉您源代码中的旧版本
    def predict(self, input_image, swap_image):
        """Run a single prediction on the model"""
        # 【新增】检查输入是否为空
        if input_image is None or swap_image is None:
            raise gr.Error("请确保同时上传了目标图片和源图片。") # 使用Gradio的方式优雅地报错
    
        try:
            # 读取图片
            target_img = cv2.imread(input_image.name)
            swap_img = cv2.imread(swap_image.name)
    
            # 分析人脸
            target_face = self.get_face(target_img)
            source_face = self.get_face(swap_img)
    
            # 【关键修正】在执行换脸前，检查是否成功找到了人脸
            if target_face is None:
                raise gr.Error("在目标图片中未能检测到人脸，请更换图片后重试。")
            if source_face is None:
                raise gr.Error("在源图片中未能检测到人脸，请更换图片后重试。")
            
            # 如果人脸都找到了，才执行核心换脸操作
            result = self.face_swapper.get(target_img, target_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:
            # 如果发生其他未知错误，也通过Gradio报错
            print(f"An unexpected error occurred: {e}")
            raise gr.Error(f"发生未知错误: {e}")


# 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,
    examples=[["input.jpg", "swap img.jpg"]])


# Launch the Gradio Interface
iface.launch()