import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np

# Load the trained Mudra model
model = YOLO("best.pt")  # path inside the Space

def predict_mudra(image):
    # Convert to RGB if needed
    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Run classification
    results = model.predict(img_rgb, task="classify", imgsz=224, verbose=False)
    result = results[0]
    class_id = result.probs.top1
    class_name = result.names[class_id]
    confidence = result.probs.top1conf.item()
    
    return f"{class_name} ({confidence*100:.1f}%)"

# Gradio interface
iface = gr.Interface(
    fn=predict_mudra,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Textbox(label="Detected Mudra"),
    title="Mudra Classifier",
    description="Upload a hand image and detect which Mudra it is."
)

iface.launch()