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

# Load the YOLOv8 model (replace with the path to your trained model if available)
model = YOLO('yolov8n.pt')  # Use a smaller model like 'yolov8n' for faster inference

# Define labels according to international standards (e.g., IEC 60617)
labels = [
    "R"  # Resistor
    "C"  # Capacitor
    "IC"  # Integrated Circuit
    "L"  # Inductor
    "Q"  # Transistor
    # Add more labels as needed
]

def predict(image):
    """Performs object detection and returns results."""
    try:
        # Convert PIL Image to NumPy array (required by Ultralytics)
        image = np.array(image) 
        results = model(image) 

        if len(results.xyxy[0]) == 0:
            print("No detections found.")  # Add debugging statement
            return "No components detected."

        detections = []
        for *xyxy, conf, cls in results.xyxy[0]:
            x1, y1, x2, y2 = map(int, xyxy)
            label = labels[int(cls)] 
            detections.append(f"{label} (Confidence: {conf:.2f})") 

        return "\n".join(detections)

    except Exception as e:
        # Handle potential errors (e.g., invalid image, model loading issues)
        print(f"Error during prediction: {e}")
        return "Error: Unable to detect components." 

# Create the Gradio interface
iface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"),  # Accepts PIL Image format
    outputs=gr.Textbox(),  # Output results as plain text
    title="PCB Component Detection",
    description="Upload an image of a PCB to identify its components.",
)

# Launch the interface
iface.launch()