Update app.py

app.py

@@ -8,15 +8,19 @@ 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
+# ✅ 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):
     """Runs YOLOv8 on the input image and returns detected defects."""
-    results = model(image, conf=0.15)  # 🔥 Lower confidence threshold to detect more defects
+    results = model(image, conf=0.15)  # 🔥 Lower confidence to detect more defects
 
     boxes = results[0].boxes.xyxy.cpu().numpy()  # Extract bounding boxes
     confs = results[0].boxes.conf.cpu().numpy()  # Extract confidence scores
@@ -26,7 +30,7 @@ def detect_pcb_faults(image):
     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)
         
-        # Get class label from model
+        # Get class label from dictionary
         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)