# Imports
import tensorflow as tf
import gradio as gr
import cv2
import numpy as np
from PIL import Image
import sys
import os
from skimage import feature

# Model path...
edges2ferrets = './e299999_generator.h5'

ferretMaker = tf.keras.models.load_model(edges2ferrets)

def transform_image(img, input_type):
    sigma = 1.7
    if input_type == "Photo":
        # Resize the input
        img = cv2.resize(img, (256, 256))

        # Convert the input photo to grayscale
        grey_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

        # Apply skimage transformation to input photo
        edges = (feature.canny(grey_img, sigma=sigma) * 255).astype(np.uint8)
        processed_input = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)

    elif input_type == "Drawing":
        # Resize the input
        img = cv2.resize(img, (256, 256))

        # Applies CV2 transformation to input drawing
        grey_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        processed_input = cv2.merge([grey_img, grey_img, grey_img])
    else:
        raise ValueError("Invalid input_type. Choose 'photo' or 'drawing'.")

    # Ferret-ing the input!
    processed_input = (processed_input - 127.5) / 127.5
    result = ferretMaker(np.expand_dims(processed_input, 0), training=True)[0].numpy()
    generated_ferret = (result * 0.5 + 0.5) * 255.0

    return generated_ferret.astype(np.uint8)

# Different canny logic is applied depending on the input the user selects. With photo, skimage is applied, and with drawing cv2 bw is applied.
iface = gr.Interface(
    fn=transform_image,
    inputs=[gr.Image(type="numpy", image_mode="RGB"), gr.Radio(["Photo", "Drawing"], label="Select Input Type")],
    outputs=gr.Image(type="numpy", image_mode="RGB"),
    live=True,
    title="Edges2Ferrets generation!",
    description="Choose between photo or drawing input to create a ferret!",
)

iface.launch(share=True)