import tensorflow as tf

import math
import numpy as np
from PIL import Image

from tensorflow.keras.preprocessing.image import img_to_array
from huggingface_hub import from_pretrained_keras
import gradio as gr

model = from_pretrained_keras("Dvjc1899/super-resolution")

def infer(image):
	img = Image.fromarray(image)
	img = img.resize((300, 300))
	ycbcr = img.convert("YCbCr")
	y, cb, cr = ycbcr.split()
	y = img_to_array(y)
	y = y.astype("float32") / 255.0

	input = np.expand_dims(y, axis=0)
	out = model.predict(input)

	out_img_y = out[0]
	out_img_y *= 255.0

	# Restore the image in RGB color space.
	out_img_y = out_img_y.clip(0, 255)
	out_img_y = out_img_y.reshape((np.shape(out_img_y)[0], np.shape(out_img_y)[1]))
	out_img_y = PIL.Image.fromarray(np.uint8(out_img_y), mode="L")
	out_img_cb = cb.resize(out_img_y.size, PIL.Image.BICUBIC)
	out_img_cr = cr.resize(out_img_y.size, PIL.Image.BICUBIC)
	out_img = PIL.Image.merge("YCbCr", (out_img_y, out_img_cb, out_img_cr)).convert(
		"RGB"
	)
	return out_img
	
iface = gr.Interface(
	fn=infer,
	title = "Super-resolution",
	inputs=gr.inputs.Image(label="Input Image"),
	outputs=[gr.outputs.Image(label="Super-resolution image"),
	],
).launch()