import gradio as gr
import cv2
import numpy as np
from huggingface_hub import hf_hub_download
import tensorflow as tf

# Load model
model = tf.keras.models.load_model(
    hf_hub_download("nharshavardhana/quickdraw_classifier", "quickdraw_classifier.keras")
)

# Class names (replace with your 50 classes)
class_names = ['anvil','banana','bowtie','butterfly','cake','carrot','cat','clock','mushroom','cup','door', 'dog','eye','fish','hexagon','moon','ice cream','pizza','umbrella','circle','star','triangle','apple', 'car', 'house', 'tree', 'cloud', 'face', 'flower', 'bird']  # Add all 50 labels

def predict_uploaded_image(img):
    # Preprocess image
    img = img.astype("float32") / 255.0
    img = 1.0 - img  # Invert colors (if needed)
    img = cv2.resize(img, (28, 28))
    img = np.expand_dims(img, axis=(0, -1))
    
    # Predict
    preds = model.predict(img)[0]
    top5 = np.argsort(preds)[::-1][:5]
    return {class_names[i]: float(preds[i]) for i in top5}

# Create a detailed UI with Blocks
with gr.Blocks(title="DoodleSense") as demo:
    gr.Markdown("# 🎨 DoodleSense")
    gr.Markdown("""
    **Draw a sketch in paint application with brush(black) of 30 px(pixels) against white background and upload the saved image** to see the top 5 predictions!  
    Try to sketch and upload any of these : 'anvil','banana','bowtie','butterfly','cake','carrot','cat','clock','mushroom','cup','door', 'dog','eye','fish','hexagon','moon','ice cream','pizza','umbrella','circle','star','triangle','apple', 'car', 'house', 'tree', 'cloud', 'face', 'flower', 'bird'.
    """)
    gr.Markdown("""
    Currently this model is trained on the [QuickDraw Dataset](https://quickdraw.withgoogle.com/data) for 30 classes.
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                image_mode="L",
            )
            gr.Examples(
                examples=["examples/butterfly.png", "examples/car.png"],  # Add your example images
                inputs=input_image,
                label="Try these examples:"
            )
        with gr.Column():
            output_label = gr.Label(num_top_classes=5, label="Top 5 Predictions")
    
    gr.Markdown("""
    ## 📖 About This Project
    - **Model**: Trained using TensorFlow/Keras on 30 QuickDraw classes.
    - **Input**: 28x28 grayscale sketches (black strokes on white background).
    - **Training Data**: 5000 samples per class from the QuickDraw dataset.
    """)

    input_image.change(predict_uploaded_image, inputs=input_image, outputs=output_label)

demo.launch(share=True)