import gradio as gr
import numpy as np
from PIL import Image
import os
os.system('pip install basicsr')
from gfpgan import GFPGANer

# installing version 1 of GFPGAN
os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth')
# installing version 1.2 of GFPGAN
os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth')
# installing version 1.3 of GFPGAN (latest)
os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth')


def interface(image: Image, model: str = "GFPGANv1.3.pth"):
    if model == "":
        model = "GFPGANv1.3.pth"
    if model != "GFPGANv1.pth" and model != "GFPGANCleanv1-NoCE-C2.pth" and model != "GFPGANv1.3.pth":
        model = "GFPGANv1.3.pth"
    restorer = GFPGANer(
        model_path=model,
        arch="original" if model == "GFPGANv1.pth" else "clean",
        bg_upsampler=None,
        channel_multiplier=1 if model == "GFPGANv1.pth" else 2,
        upscale=2)
    img = np.array(image)[:, :, ::-1].copy()
    cropped_faces, restored_faces, restored_img = restorer.enhance(
        img,
        align=False,
        only_center_face=False,
    )
    return restored_img


gr.Interface(
    interface,
    [
        gr.components.Image(
            type="pil",
            label="Image",
        ),
        gr.components.Radio([
            "GFPGANv1.pth",
            "GFPGANCleanv1-NoCE-C2.pth",
            "GFPGANv1.3.pth",
        ],
                            label="model",
                            default="GFPGANv1.3.pth",
                            type="value")
    ],
    [gr.components.Image(label="Enhanced Image")],
).launch()