import gradio as gr
from PIL import Image

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan.utils import RealESRGANer

# model load
netscale = 4
super_res_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
super_res_upsampler = RealESRGANer(scale=netscale, model_path='model_zoo/RealESRGAN_x4plus.pth', model=super_res_model, tile=0, 
                                    tile_pad=10, pre_pad=0, half=False, gpu_id=None)
fisheye_correction_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
fisheye_correction_upsampler = RealESRGANer(scale=netscale, model_path='model_zoo/RealESRGAN_x4plus_fine_tuned_400k.pth', model=fisheye_correction_model, tile=0, 
                                    tile_pad=10, pre_pad=0, half=False, gpu_id=None)

def predict(radio_btn, input_img):
    out = None

    # preprocess input
    if(input_img is not None):
        if(radio_btn == 'Super resolution'):
            upsampler = super_res_upsampler
        else:
            upsampler = fisheye_correction_upsampler
        output, _ = upsampler.enhance(input_img, outscale=4)

        # convert to pil image
        out = Image.fromarray(output)
    return out


gr.Interface(
    fn=predict,
    inputs=[
        gr.Radio(choices=["Super resolution", "Distortion correction"], value="Super resolution", label="Select task:"), gr.inputs.Image()
    ],
    outputs=[
        gr.inputs.Image()
    ],
    title="Real-ESRGAN moon distortion",
    description="Description of the app",
    examples=[
        ["Super resolution", "render0001.png"], ["Super resolution", "render1546.png"], ["Super resolution", "render1682.png"], 
        ["Distortion correction", "render0001_DC.png"], ["Distortion correction", "render1546_DC.png"], ["Distortion correction", "render1682_DC.png"]
    ]
).launch()