# -*- coding: utf-8 -*-
import numpy as np

import tensorflow as tf

import gradio as gr

from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras("keras-io/CycleGAN", compile=False)

# Define the standard image size.
orig_img_size = (286, 286)
# Size of the random crops to be used during training.
input_img_size = (256, 256, 3)


def normalize_img(img):
    img = tf.cast(img, dtype=tf.float32)
    # Map values in the range [-1, 1]
    return (img / 127.5) - 1.0


def preprocess_test_image(img):
    # Only resizing and normalization for the test images.
    img = tf.image.resize(img, [input_img_size[0], input_img_size[1]])
    img = normalize_img(img)
    return img


# img_path = './n02381460_1010.jpg'


def generate_img(img_path):
    img = tf.io.read_file(img_path)
    img = tf.image.decode_png(img)
    img = tf.expand_dims(img, axis=0)
    img = preprocess_test_image(img)
    prediction = model(img, training=False)[0].numpy()
    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)
    return prediction


image = gr.inputs.Image(type="filepath")
op = gr.outputs.Image(type="numpy")

iface = gr.Interface(
    generate_img,
    image,
    op,
    title="CycleGAN",
    description='Keras Implementation of CycleGAN model using Horse to Zebra dataset',
    article='Author: <a href="https://huggingface.co/anuragshas">Anurag Singh</a>. Based on the keras example from <a href="https://keras.io/examples/generative/cyclegan/">A_K_Nain</a>',
    examples=["n02381460_360.jpg", "n02381460_4410.jpg"],
)

iface.launch(cache_examples=True)