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

# ✅ Load YOLOv8 model from Hugging Face
model_path = "https://huggingface.co/Sakthi3214/pcb_detection/resolve/main/best.pt"
model = YOLO(model_path)

# ✅ Define class names (Manually if model.names is empty)
class_names = {
    0: "Missing Hole",
    1: "Mouse Bite",
    2: "Open Circuit",
    3: "Short",
    4: "Spur",
    5: "Copper",
} if not model.names else model.names  # Use model.names if available

def detect_pcb_faults(image, conf_threshold):
    """Runs YOLOv8 on the input image with an adjustable confidence threshold."""
    results = model(image, conf=conf_threshold)[0]  # 🔥 Adjustable confidence

    boxes = results.boxes.xyxy.cpu().numpy()  # Extract bounding boxes
    confs = results.boxes.conf.cpu().numpy()  # Extract confidence scores
    class_ids = results.boxes.cls.cpu().numpy()  # Extract class IDs

    # ✅ Draw bounding boxes and labels
    for (x1, y1, x2, y2), conf, class_id in zip(boxes, confs, class_ids):
        cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 3)  # 🟢 Thicker Box
        
        # Get class label from dictionary
        label = f"{class_names.get(int(class_id), 'Unknown')} ({conf:.2f})"
        
        # 🔥 Larger Text Size & Thicker Font
        cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 3)

    return image

# ✅ Gradio UI with Confidence Threshold Slider
gr.Interface(
    fn=detect_pcb_faults,
    inputs=[
        gr.Image(type="numpy"),
        gr.Slider(minimum=0.1, maximum=0.9, value=0.18, label="Confidence Threshold")
    ],
    outputs=gr.Image(type="numpy"),
    title="PCB Fault Detection",
    description="Upload a PCB image to detect defects using YOLOv8. Adjust the confidence threshold to fine-tune detections."
).launch()