import os
import cv2
import numpy as np
import torch
from ultralytics import YOLO
import gradio as gr
import traceback
# -----------------------------
# 1. YOLO model path
# -----------------------------
YOLO_MODEL_PATH = "best.pt"  # Push this small model to HF repo

# -----------------------------
# 2. Load YOLO model
# -----------------------------
yolo_model = YOLO(YOLO_MODEL_PATH)
yolo_model.eval()

# -----------------------------
# 3. Inference function
# -----------------------------
def predict_asl(image):
    try:
        if image is None:
            raise ValueError("No image uploaded")
        
        img = image.copy()
        h, w, _ = img.shape
        print(f"🔹 Uploaded image shape: {img.shape}, dtype: {img.dtype}")

        # --- YOLO prediction directly on NumPy array ---
        results = yolo_model.predict(img, imgsz=300, verbose=False)[0]
        pred_idx = results.probs.top1
        pred_label = results.names[pred_idx]
        confidence = results.probs.data[pred_idx].item()

        # Overlay prediction text on original image
        cv2.putText(
            img,
            f"{pred_label} ({confidence:.2f})",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 0, 255),
            2,
            cv2.LINE_AA
        )

        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), pred_label, round(confidence, 2)

    except Exception as e:
        print("❌ Error in predict_asl:", e)
        traceback.print_exc()
        return image, "Error", 0.0

# -----------------------------
# 4. Gradio Interface
# -----------------------------
title = "🖐️ ASL Letter Classifier"
description = "Upload a hand sign image and see the predicted letter and confidence."

iface = gr.Interface(
    fn=predict_asl,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(type="numpy", label="Original Image with Prediction"),
        gr.Textbox(label="Predicted Letter"),
        gr.Textbox(label="Confidence")
    ],
    title=title,
    description=description,
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch(share=True)