import tensorflow as tf
import gradio as gr
from scipy.ndimage import zoom
import cv2
import numpy as np 

# Load the pre-trained model
from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras("manufy/mnist_model_keras")
#model = tf.keras.models.load_model('mnist_model.keras')

def recognize_digit(input):
    image = input['layers'][0]
    print("Type of image variable:", type(image))
    dimensions = image.shape
    print("Dimensions of the image:", dimensions)
    # Resize image to 28x28 using interpolation
    resized_image = cv2.resize(image, (28, 28))
     # Convert image to grayscale
    gray_image = cv2.cvtColor(resized_image, cv2.COLOR_RGBA2GRAY)
    # Reshape the image to add a single channel dimension
    final_image = np.expand_dims(gray_image, axis=-1)
    
    print("Final Dimensions of the image:", final_image.shape)
   
    #if image.shape[-1] == 3:  # If the image has 3 color channels (RGB)
    #    image = image.mean(axis=-1)  # Convert to grayscale
    # Resize the image to 28x28
    #image = zoom(image, (28 / image.shape[0], 28 / image.shape[1]))
    final_image = final_image.reshape((1, 28, 28, 1)).astype('float32') / 255
    prediction = model.predict(final_image)
    return {str(i): float(prediction[0][i]) for i in range(10)}, input['layers'][0]
#final_image[0, :, :, 0]

    #return ''

# Configure the Sketchpad
sketchpad = gr.Sketchpad()

# Create the Gradio interface
interface = gr.Interface(
    fn=recognize_digit, 
    inputs="sketchpad",
    outputs=[gr.Label(num_top_classes=3), "image"],
    live=True,
    title="Digit Recognizer",
    description="Draw a digit and let the model recognize it."
)

# Launch the interface
interface.launch()