from PIL import Image
import numpy as np
import gradio as gr
from tensorflow.keras.models import load_model

model = load_model('mio_modello.h5')

# Image processing
def preprocess_image(image):
    
    image = Image.fromarray(image)  
    image = image.convert("RGB")
    image = image.resize((64, 64))
    image_array = np.array(image) / 255.0  
    
    image_array = np.expand_dims(image_array, axis=0)  # (1, 64, 64, 3)
    
    return image_array

def classify_image(image):
    image_array = preprocess_image(image)
    
    prediction = model.predict(image_array)[0] 
    
    print(f"Raw model predictions: {prediction}")
    
    predicted_class_idx = np.argmax(prediction)
    
    class_labels = ['Chihuahua', 'Muffin']
    
    print(f"Predicted class: {class_labels[predicted_class_idx]} with confidence {prediction[predicted_class_idx]}")
    
    confidence_scores = {class_labels[i]: float(prediction[i]) for i in range(len(class_labels))}
    
    return confidence_scores

examples = ['chihuahua1.jpg', 'chihuahua2.jpg', 'chihuahua3.jpg', 'muffin1.jpg', 'muffin2.jpg']

#  Gradio Interface
gr.Interface(
    fn=classify_image,  
    inputs=gr.Image(type="numpy"), 
    outputs=gr.Label(num_top_classes=2), 
    examples=examples,
    title="Chihuahua vs Muffin Classifier",
    description="don't know if you should cuddle it or eat it? find out uploading a picture. PS: no muffin and no chihuahua got hurt in this project"
).launch()