fix off mapping

app.py

@@ -24,10 +24,9 @@ DETECTION_MODELS = [
     "dinov3/convnext-tiny-ltdetr-coco"
 ]
 
-# UPDATED: Added Base (vitb16) and Large (vitl16) for better accuracy
 SEGMENTATION_MODELS = [
-    "dinov3/vitb16-eomt-coco",       # Base (Recommended Balance)
-    "dinov3/vitl16-eomt-coco",       # Large (Best Accuracy, Slower)
+    "dinov3/vitb16-eomt-coco",       # Base (Balanced)
+    "dinov3/vitl16-eomt-coco",       # Large (Best Accuracy)
     "dinov3/vits16-eomt-coco"        # Small (Fastest)
 ]
 
@@ -49,8 +48,10 @@ COCO_DETECTION_CLASSES = [
     "scissors", "teddy bear", "hair drier", "toothbrush"
 ]
 
-# COCO-Stuff (171 Classes) - Standard Mapping
+# COCO-Stuff (171 Classes)
+# FIX: Added 'unlabeled' at index 0 so 'person' aligns with Index 1
 COCO_STUFF_CLASSES = [
+    "unlabeled", # Index 0 (Background)
     "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
     "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
     "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
@@ -163,10 +164,9 @@ def run_segmentation(model, image_input, original_image):
     labels_to_draw = []
     
     for cls_id in unique_classes:
-        # Safety check: skip 'background' class
-        if cls_id == 255 or cls_id == -1: continue
-
-        # COCO-Stuff mapping
+        # 0 is usually background/unlabeled in this list
+        if cls_id == 0: continue
+        # Safety check for range
         if cls_id < 0 or cls_id >= len(current_classes): continue
         
         class_name = current_classes[cls_id]
@@ -177,6 +177,7 @@ def run_segmentation(model, image_input, original_image):
         colored_mask[mask_np == cls_id] = color
 
         y_indices, x_indices = np.where(mask_np == cls_id)
+        # Filter small noise
         if len(y_indices) > 200: 
             centroid_y = int(np.mean(y_indices))
             centroid_x = int(np.mean(x_indices))
@@ -208,6 +209,7 @@ with gr.Blocks(theme=theme) as demo:
             
             with gr.Accordion("Settings", open=True):
                 conf_slider = gr.Slider(0.0, 1.0, value=0.4, step=0.05, label="Confidence (Detection Only)")
+                # BUMPED DEFAULT TO 640 for sharper masks
                 res_slider = gr.Slider(384, 1024, value=640, step=32, label="Inference Resolution")
                 
                 model_selector = gr.Dropdown(
@@ -236,7 +238,7 @@ with gr.Blocks(theme=theme) as demo:
         examples=[
             ["http://farm3.staticflickr.com/2547/3933456087_6a4dfb4736_z.jpg", 0.4, 640, DEFAULT_MODEL],
             ["https://farm3.staticflickr.com/2294/2193565429_aed7c9ff98_z.jpg", 0.4, 640, DEFAULT_MODEL],
-            ["http://farm9.staticflickr.com/8092/8400332884_102a62b6c6_z.jpg", 0.4, 512, "dinov3/vitl16-eomt-coco"],
+            ["http://farm9.staticflickr.com/8092/8400332884_102a62b6c6_z.jpg", 0.4, 640, "dinov3/vitn16-eomt-coco"],
            ],
         outputs=[output_img, output_text, output_json],
         fn=run_prediction,