import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO

# Load YOLOv11 Model
model_path = "best.pt" 
model = YOLO(model_path)

def predict(image):
    image = np.array(image)
    results = model(image, conf=0.85)

    detected_classes = set()  # Track unique detected classes
    labels = []

    # Draw bounding boxes and extract labels
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            cls = int(box.cls[0])
            class_name = model.names[cls]

            detected_classes.add(class_name)  # Store detected class
            label = f"{class_name} {conf:.2f}"
            labels.append(label)

            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # Define possible classes (adjust based on your dataset)
    possible_classes = {"front", "back"}

    # Identify missing class if any
    missing_classes = possible_classes - detected_classes
    if missing_classes:
        labels.append(f"Missing: {', '.join(missing_classes)}")

    return Image.fromarray(image), labels

# Gradio Interface
iface = gr.Interface(
    fn=predict, 
    inputs="image", 
    outputs=["image", "text"],  # Returning both image and detected labels
    title="YOLOv11 Object Detection (Front & Back Card)"
)

iface.launch()