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

# Load YOLOv8 model
model_path = "best(1).pt"
model = YOLO(model_path)

# Assign unique colors to each class
colors = {name: (np.random.randint(100, 255), np.random.randint(100, 255), np.random.randint(100, 255)) for name in model.names.values()}

# Function to detect components in a PCB image
def detect_pcb(image, conf_threshold=0.5):
    results = model(image)
    detected_labels = []
    
    # Convert image format
    image = np.array(image, dtype=np.uint8)
    
    # Process results
    for result in results:
        for box in result.boxes:
            conf = box.conf[0].item()  # Confidence score
            if conf < conf_threshold:
                continue  # Ignore low-confidence detections

            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates
            cls = int(box.cls[0])  # Class index
            label = model.names[cls]  # Component name
            detected_labels.append(f"{label} ({conf:.2f}) at ({x1}, {y1})")

            # Define colors and font
            color = colors[label]
            font = cv2.FONT_HERSHEY_SIMPLEX
            
            # Draw bounding box
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 3)

            # Add text with a semi-transparent background
            text = f"{label} ({conf:.2f})"
            (w, h), _ = cv2.getTextSize(text, font, 0.7, 2)
            cv2.rectangle(image, (x1, y1 - h - 10), (x1 + w + 4, y1), color, -1)  # Background
            cv2.putText(image, text, (x1 + 2, y1 - 5), font, 0.7, (255, 255, 255), 2)

    return image, "\n".join(detected_labels)

# Create Gradio Interface
iface = gr.Interface(
    fn=detect_pcb,
    inputs=[gr.Image(type="numpy"), gr.Slider(0.1, 1.0, value=0.5, label="Confidence Threshold")],
    outputs=[gr.Image(type="numpy"), gr.Textbox(label="Detected Components")],
    title="🔍 PCB Component Detection",
    description="Upload a PCB image to detect components with a clearer output. Adjust confidence threshold to filter results.",
)

iface.launch()