app.py

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()