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

# -----------------------------
# 1. YOLO model
# -----------------------------
YOLO_MODEL_PATH = "best.pt"
yolo_model = YOLO(YOLO_MODEL_PATH)


# -----------------------------
# 2. Reference alphabet image (WebP)
# -----------------------------
REFERENCE_IMAGE_PATH = "asl_alphabet.jpg"
reference_img = cv2.imread(REFERENCE_IMAGE_PATH)
reference_img = cv2.cvtColor(reference_img, cv2.COLOR_BGR2RGB)

# -----------------------------
# 3. Prediction function
# -----------------------------
def predict_asl(image):
    try:
        if image is None:
            raise ValueError("No image uploaded")
        
        img = image.copy()

        results = yolo_model.predict(image, imgsz=300, verbose=False)
        pred_idx = results[0].probs.top1
        pred_label = results[0].names[pred_idx]
        confidence = results[0].probs.top1conf.item() 


        # Overlay prediction text
        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 Layout
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🖐️ ASL Letter Classifier")
    gr.Markdown("Upload a hand sign image and see the predicted letter and confidence. The full ASL alphabet is always shown on the left as a reference.")

    with gr.Row():
        # Left column: reference alphabet
        with gr.Column(scale=1):
            gr.Image(value=reference_img, type="numpy", label="ASL Alphabet Reference")
        
        # Right column: upload & prediction
        with gr.Column(scale=2):
            input_image = gr.Image(type="numpy", label="Upload your ASL Letter")
            output_image = gr.Image(type="numpy", label="Prediction")
            pred_text = gr.Textbox(label="Predicted Letter")
            confidence_text = gr.Textbox(label="Confidence")
            
            input_image.change(
                fn=predict_asl,
                inputs=input_image,
                outputs=[output_image, pred_text, confidence_text]
            )

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