import gradio as gr
import tensorflow as tf
from PIL import Image, ImageOps
import numpy as np
import os
import cv2

model_path = "IM_417_128.keras"

print("Loading Model...")
model = tf.keras.models.load_model(model_path)
model.summary()

labelled_images = [f"labelled/{i}.png" for i in range(1, 418)]

# Define regression function
def predict_regression(image):
    # Preprocess image

    print(f" {type(image)} image layers ===> { len(image['layers']) }")
    image = np.squeeze(image["layers"][0][:, :, -1:], axis=-1)
    
    image = np.clip(image, 0, 255).astype('uint8')

    resized_image = cv2.resize(image, (128, 128), interpolation=cv2.INTER_CUBIC)
    
    # Convert to grayscale
    # gray_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)

    # Create a mask for values greater than 10
    mask = resized_image > 10
    
    # Use the mask to assign values
    gray_image = np.where(mask, 255, 0)

    print(f"gray_image image TEST :: {gray_image.shape}")
    
    # Convert the image to a NumPy array
    image_array = np.array(gray_image)
    # Invert the colors
    inverted_image_array = 255 - image_array
    
    # Predict
    predictions = model.predict(inverted_image_array[None, ...])  # Assuming single regression value
    predicted_output = np.argmax(predictions, axis=1)
    predicted_label = f"labelled/{str(predicted_output[0])}.png"
    print(f"prediction IM-417 character :: {predicted_label} :: source {predicted_output}")
    output_data = cv2.imread(predicted_label, cv2.COLOR_BGR2GRAY)
    return cv2.resize(output_data, (128, 128)) , image_array

# Create Gradio interface
output_image = gr.Image(height=150, width=150)
i_0 = gr.Image(height=150, width=150)

# sketchpad = gr.Sketchpad(flatten=True) 
interface = gr.Interface(fn=predict_regression,
             inputs="sketchpad",
             outputs=[output_image, i_0],
             live=True,
             description="A simple mlp classification model for image classification using the sign list from 'The Indus Scripts: Texts, Concordance and Tables'"
)


interface.launch()