Update app.py

app.py

@@ -8,24 +8,35 @@ import gradio as gr
 model_path = "https://huggingface.co/Sakthi3214/pcb_detection/resolve/main/best.pt"
 model = YOLO(model_path)
 
+# ✅ Check if model has class names
+if model.names:
+    class_names = model.names
+else:
+    class_names = {0: "Defect"}  # Default label if no names are found
+
 def detect_pcb_faults(image):
-    """Runs YOLOv8 on the input image and returns detected objects."""
-    results = model(image)  # Run inference
+    """Runs YOLOv8 on the input image and returns detected defects."""
+    results = model(image, conf=0.15)  # 🔥 Lower confidence threshold to detect more defects
+
     boxes = results[0].boxes.xyxy.cpu().numpy()  # Extract bounding boxes
     confs = results[0].boxes.conf.cpu().numpy()  # Extract confidence scores
+    class_ids = results[0].boxes.cls.cpu().numpy()  # Extract class IDs
 
-    # Draw bounding boxes
-    for (x1, y1, x2, y2), conf in zip(boxes, confs):
+    # ✅ 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), 2)
-        cv2.putText(image, f"{conf:.2f}", (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        
+        # Get class label from model
+        label = f"{class_names.get(int(class_id), 'Unknown')} ({conf:.2f})"
+        cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
 
     return image
 
-# ✅ Gradio UI
+# ✅ Gradio UI for PCB Fault Detection
 gr.Interface(
     fn=detect_pcb_faults,
     inputs=gr.Image(type="numpy"),
     outputs=gr.Image(type="numpy"),
     title="PCB Fault Detection",
-    description="Upload a PCB image to detect defects using YOLOv8."
+    description="Upload a PCB image to detect defects using YOLOv8. The model will highlight and label detected faults."
 ).launch()