Upload app.py

app.py

@@ -31,7 +31,7 @@ def draw_boxes(img_rgb, boxes, labels, color=(0, 220, 100)):
     return out
 
 # ─── Model Functions ──────────────────────────────────────────────────────────
-def run_yolo_generic(img_rgb, model_path, target_class, color):
+def run_yolo_generic(img_rgb, model_path, target_classes, color):
     from ultralytics import YOLO
     t0 = time.time()
     model = YOLO(model_path)
@@ -50,13 +50,13 @@ def run_yolo_generic(img_rgb, model_path, target_class, color):
             result.boxes.cls,
             result.boxes.conf
         ):
-            if int(cls) != target_class:
+            if int(cls) not in target_classes:
                 continue
             mask_np = mask.cpu().numpy()
             mask_r  = cv2.resize(mask_np, (w, h), interpolation=cv2.INTER_NEAREST)
             combined_mask |= mask_r > 0.5
             boxes.append(box.cpu().tolist())
-            class_name = "mirror" if target_class == 0 else "car"
+            class_name = "mirror" if 0 in target_classes or 1 in target_classes else "car"
             labels.append(f"{class_name} {conf:.2f}")
 
     out = apply_mask_overlay(img_rgb, combined_mask, color=color)
@@ -64,7 +64,7 @@ def run_yolo_generic(img_rgb, model_path, target_class, color):
     n = len(boxes)
     return out, f"Found: {n} | Inference Time: {elapsed:.2f}s"
 
-def run_sam_generic(img_rgb, yolo_model_path, target_class, color):
+def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
     try:
         from segment_anything import sam_model_registry, SamPredictor
         import urllib.request, os
@@ -95,7 +95,7 @@ def run_sam_generic(img_rgb, yolo_model_path, target_class, color):
                 yolo_res.boxes.cls,
                 yolo_res.boxes.conf
             ):
-                if int(cls) != target_class:
+                if int(cls) not in target_classes:
                     continue
                 box_np = box.cpu().numpy()
                 masks_sam, _, _ = predictor.predict(
@@ -104,7 +104,7 @@ def run_sam_generic(img_rgb, yolo_model_path, target_class, color):
                 )
                 combined_mask |= masks_sam[0]
                 boxes_list.append(box_np.tolist())
-                class_name = "mirror" if target_class == 0 else "car"
+                class_name = "mirror" if 0 in target_classes or 1 in target_classes else "car"
                 labels.append(f"{class_name} {conf:.2f}")
 
         elapsed = time.time() - t0
@@ -204,16 +204,16 @@ def process_image(img_rgb, model_name):
         
     try:
         if model_name == "YOLOv8x-seg (Custom Mirror)":
-            return run_yolo_generic(img_rgb, "best.pt", target_class=0, color=(50, 220, 100))
+            return run_yolo_generic(img_rgb, "best.pt", target_classes=[0, 1], color=(50, 220, 100))
             
         elif model_name == "YOLOv8x (Pretrained Car)":
-            return run_yolo_generic(img_rgb, "yolov8x-seg.pt", target_class=2, color=(0, 200, 255))
+            return run_yolo_generic(img_rgb, "yolov8x-seg.pt", target_classes=[2], color=(0, 200, 255))
             
         elif model_name == "SAM + YOLO (Custom Mirror)":
-            return run_sam_generic(img_rgb, "best.pt", target_class=0, color=(255, 80, 160))
+            return run_sam_generic(img_rgb, "best.pt", target_classes=[0, 1], color=(255, 80, 160))
             
         elif model_name == "SAM + YOLO (Pretrained Car)":
-            return run_sam_generic(img_rgb, "yolov8x-seg.pt", target_class=2, color=(200, 80, 255))
+            return run_sam_generic(img_rgb, "yolov8x-seg.pt", target_classes=[2], color=(200, 80, 255))
             
         elif model_name == "Mask R-CNN (Pretrained Car)":
             return run_maskrcnn(img_rgb)