import gradio as gr
import tensorflow as tf
import numpy as np
import json
from PIL import Image, ImageOps  # Added ImageOps for inversion

# 1. Load Model and Labels
model = tf.keras.models.load_model('devanagari_model.keras')

with open('labels.json', 'r') as f:
    labels = json.load(f)

# 2. Preprocessing Function
def process_image(image):
    # Convert to grayscale (L)
    image = image.convert('L')
    
    # --- CRITICAL FIX START ---
    # Invert colors: Black text/White bg -> White text/Black bg
    # This matches the UCI dataset format used in training.
    image = ImageOps.invert(image)
    # --- CRITICAL FIX END ---
    
    # Resize to 32x32 (dataset size)
    image = image.resize((32, 32))
    
    # Convert to array
    img_array = np.array(image)
    
    # Normalize to 0-1
    img_array = img_array / 255.0
    
    # Add batch dimension (1, 32, 32, 1)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = np.expand_dims(img_array, axis=-1)
    
    return img_array

# 3. Prediction Function
def predict_character(image):
    if image is None:
        return "Please upload an image."
    
    processed_img = process_image(image)
    predictions = model.predict(processed_img)
    
    # Get top prediction
    predicted_class_index = np.argmax(predictions)
    
    # JSON keys are strings, so cast index to str
    predicted_label = labels[str(predicted_class_index)]
    
    # Convert numpy float to python float for Gradio
    confidence = float(np.max(predictions))
    
    # Return dictionary for Gradio Label output
    return {predicted_label: confidence}

# 4. Gradio Interface
iface = gr.Interface(
    fn=predict_character,
    inputs=gr.Image(type="pil", label="Upload Character Image"),
    outputs=gr.Label(num_top_classes=3),
    title="Devanagari Character Recognition (Lightweight)",
    description="Upload a handwritten Hindi/Devanagari character. This model is optimized for low-resource environments."
)

iface.launch()