import cv2
import os
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

import gradio as gr
import torchvision.transforms as transforms


model_path = "Trained_ESRGAN.pth"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model)



def esrgan(input_image):
  output_img, _ = upsampler.enhance(input_image, outscale=3.5)
  filename = "output.jpg"
  output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB)
  cv2.imwrite(filename, output_img)    
  return filename


# Define the Gradio app interface
inputs = gr.Image(label="Input Image")
outputs = gr.Image(label="Enhanced_Image.")
title = "Image Super-Resolution Using ESR-GAN"
description = "Enhance the Quality of your Low Resolution Images To High Resolution Using Artificial Intelligence"

iface = gr.Interface(fn=esrgan, inputs=inputs, outputs=outputs, title=title, description=description, allow_flagging="never")

iface.launch(inline = False)