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

import gradio as gr
from skimage.restoration import inpaint as p
import torchvision.transforms as transforms


def restore_image(input_image):
    img = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY)
    threshold_value = 200 
    _, mask = cv2.threshold(img, threshold_value, 255, cv2.THRESH_BINARY)
    channels = cv2.split(img)
    inpaint_channels = []
    for channel in channels:
        inpaint_result = p.inpaint_biharmonic(channel, mask)
        inpaint_channels.append(inpaint_result)

    result_img = cv2.merge(inpaint_channels)
    filename = "output.jpg"
    cv2.imwrite(filename, result_img)    
    return filename


# Define the Gradio app interface
inputs = gr.Image(label="Upload Image")
outputs = gr.Image(label="Restored_Image.")
title = "Image Restoration Using Pix2Pix-GAN"
description = "Restore the Quality of your Old damaged Images To New Looking Images Using Artificial Intelligence"

iface = gr.Interface(fn=restore_image, inputs=inputs, outputs=outputs, title=title, description=description, allow_flagging="never")
iface.launch(inline = False)