File size: 1,673 Bytes
8c1e0a0 9e7454a d207917 8c1e0a0 11db117 8c1e0a0 d207917 8c1e0a0 59322b3 4e1900f 59322b3 d207917 8c1e0a0 d207917 8c1e0a0 11db117 0e263e9 8c1e0a0 11db117 e1b441b 11db117 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | 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()
|