fix

app.py

@@ -53,7 +53,7 @@ def prepare_masks(filtered_annotations):
         segmentation = annotation["segmentation"]
         mask_id = annotation["id"]
         decoded_mask = decode_segmentation(segmentation)  # Decode binary mask
-        masks.append((decoded_mask, f"mask {idx + 1}"))  # Add mask and its label
+        masks.append((decoded_mask, f"{mask_id}"))  # Add mask and its label
     return masks
 
 # Load image and annotations dynamically
@@ -80,24 +80,20 @@ def load_image_and_masks(image_id):
     return image_path, masks
 
 # Gradio event function to display captions
-# def display_caption(evt: gr.SelectData, masks, image_id):
-#     # Extract the mask ID from the label
-#     mask_id = int(masks[evt.index][1])  # Get the label corresponding to the selected mask
-#     caption_data = captions_data.get(image_id, {}).get(str(mask_id), "No caption found")
-#     return caption_data
-
 def display_caption(evt: gr.SelectData, masks, image_id):
-    # Get the label, e.g., "mask 1"
-    label = masks[evt.index][1]
-    # Extract the numeric part: split by "mask" and strip whitespace
-    try:
-        mask_number = int(label.split("mask")[1].strip())
-    except (IndexError, ValueError):
-        return "Invalid mask label"
-    
-    caption_data = captions_data.get(image_id, {}).get(str(mask_number), "No caption found")
+    # Extract the mask ID from the label
+    # mask_id = int(masks[evt.index][1])  # Get the label corresponding to the selected mask
+    mask_id = masks[evt.index][1]  # Get the label corresponding to the selected mask
+    caption_data = captions_data.get(image_id, {}).get(str(mask_id), "No caption found")
     return caption_data
 
+# def display_caption(evt: gr.SelectData, masks, image_id):
+#     # Get the label, e.g., "mask 1"
+#     label = masks[evt.index][1]
+#     # Use the label directly as key in the captions_data
+#     caption_data = captions_data.get(image_id, {}).get(label, "No caption found")
+#     return caption_data
+
 # Gradio event function to update image
 def update_image(image_index):
     image_id = image_ids[image_index]