change segm to coco stuff

app.py

@@ -154,7 +154,8 @@ def run_segmentation(model, image_input, original_image):
     mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
     mask_np = cv2.resize(mask_np, original_image.size, interpolation=cv2.INTER_NEAREST)
 
-    current_classes = COCO_STUFF_CLASSES
+    #current_classes = COCO_STUFF_CLASSES
+    current_classes = ["unlabeled"] + COCO_DETECTION_CLASSES
 
     h, w = mask_np.shape
     colored_mask = np.zeros((h, w, 3), dtype=np.uint8)
@@ -164,9 +165,9 @@ def run_segmentation(model, image_input, original_image):
     labels_to_draw = []
     
     for cls_id in unique_classes:
-        # 0 is usually background/unlabeled in this list
-        if cls_id == 0: continue
-        # Safety check for range
+        # Skip background (Index 0)
+        if cls_id == 0 or cls_id == 255: continue
+
         if cls_id < 0 or cls_id >= len(current_classes): continue
         
         class_name = current_classes[cls_id]
@@ -190,8 +191,9 @@ def run_segmentation(model, image_input, original_image):
         cv2.putText(blended, text, (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3, cv2.LINE_AA)
         cv2.putText(blended, text, (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
     
-    analytics_text = f"Scene Contains (Segmentation):\n" + (", ".join(sorted(list(found_classes))) if found_classes else "None")
-    
+    analytics_text = f"Scene Contains (COCO Objects):\n" + (", ".join(sorted(list(found_classes))) if found_classes else "None")
+
+
     return Image.fromarray(blended), analytics_text, {"classes_found": list(found_classes)}
 
 # --- GRADIO UI ---