import os
import time
from PIL import Image
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import matplotlib.pyplot as plt
import gradio as gr

# Declaring Constants
SAVED_MODEL_PATH = "https://tfhub.dev/captain-pool/esrgan-tf2/1"

def resize(width,img):
  basewidth = width
  img = Image.open(img)
  wpercent = (basewidth/float(img.size[0]))
  hsize = int((float(img.size[1])*float(wpercent)))
  img = img.resize((basewidth,hsize), Image.ANTIALIAS)
  img.save('somepic.jpg')
  return 'somepic.jpg'
  
def preprocess_image(image_path):
  """ Loads image from path and preprocesses to make it model ready
      Args:
        image_path: Path to the image file
  """
  hr_image = tf.image.decode_image(tf.io.read_file(image_path))
  # If PNG, remove the alpha channel. The model only supports
  # images with 3 color channels.
  if hr_image.shape[-1] == 4:
    hr_image = hr_image[...,:-1]
  hr_size = (tf.convert_to_tensor(hr_image.shape[:-1]) // 4) * 4
  hr_image = tf.image.crop_to_bounding_box(hr_image, 0, 0, hr_size[0], hr_size[1])
  hr_image = tf.cast(hr_image, tf.float32)
  return tf.expand_dims(hr_image, 0)

  
def plot_image(image):
  """
    Plots images from image tensors.
    Args:
      image: 3D image tensor. [height, width, channels].
      title: Title to display in the plot.
  """
  image = np.asarray(image)
  image = tf.clip_by_value(image, 0, 255)
  image = Image.fromarray(tf.cast(image, tf.uint8).numpy())
  return image
  
model = hub.load(SAVED_MODEL_PATH)
def inference(img):
  resize_image = resize(256,img)
  hr_image = preprocess_image(resize_image)
  fake_image = model(hr_image)
  fake_image = tf.squeeze(fake_image)
  pil_image = plot_image(tf.squeeze(fake_image))
  return pil_image
  
title="esrgan-tf2"
description="Enhanced Super Resolution GAN for image super resolution. Produces x4 Super Resolution Image from images of {Height, Width} >=64. Works best on Bicubically downsampled images. (*This is because, the model is originally trained on Bicubically Downsampled DIV2K Dataset*)"
article = "<p style='text-align: center'><a href='https://tfhub.dev/captain-pool/esrgan-tf2/1' target='_blank'>Tensorflow Hub</a></p>"
examples=[['input.png']]
gr.Interface(inference,gr.inputs.Image(type="filepath"),"image",title=title,description=description,article=article,examples=examples).launch(enable_queue=True)