Update app.py

app.py

@@ -8,7 +8,10 @@ from rdp import rdp
 dataset = load_dataset("dwb2023/brain-tumor-image-dataset-semantic-segmentation", split="test")
 
 def simplify_segmentation(segmentation, max_points=20):
-    simplified = rdp(np.array(segmentation), epsilon=1.0)
+    if not segmentation:
+        return []
+    epsilon = 1.0
+    simplified = rdp(np.array(segmentation), epsilon=epsilon)
     while len(simplified) > max_points:
         epsilon *= 1.5
         simplified = rdp(np.array(segmentation), epsilon=epsilon)
@@ -30,15 +33,23 @@ def draw_annotations(index):
         bbox = record["bbox"]
         draw.rectangle([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]], outline="red", width=2)
 
-        # Draw original segmentation mask
+        # Draw segmentation mask
         segmentation = record["segmentation"]
-        for seg in segmentation:
-            draw.polygon(seg, outline="blue", width=2)
+        original_points = 0
+        simplified_points = 0
+        
+        if segmentation and len(segmentation) > 0 and len(segmentation[0]) > 0:
+            for seg in segmentation:
+                if len(seg) > 0:
+                    draw.polygon(seg, outline="blue", width=2)
+                    original_points += len(seg)
 
-        # Simplify and draw simplified segmentation mask
-        simplified_segmentation = [simplify_segmentation(seg) for seg in segmentation]
-        for seg in simplified_segmentation:
-            draw.polygon(seg, outline="green", width=2)
+            # Simplify and draw simplified segmentation mask
+            simplified_segmentation = [simplify_segmentation(seg) for seg in segmentation if len(seg) > 0]
+            for seg in simplified_segmentation:
+                if len(seg) > 0:
+                    draw.polygon(seg, outline="green", width=2)
+                    simplified_points += len(seg)
 
         category_id = record["category_id"]
         area = record["area"]
@@ -48,8 +59,8 @@ def draw_annotations(index):
         info += f"Image ID: {record['id']}\n"
         info += f"Category ID: {category_id}\n"
         info += f"Bounding Box: [{bbox[0]:.2f}, {bbox[1]:.2f}, {bbox[2]:.2f}, {bbox[3]:.2f}]\n"
-        info += f"Original Segmentation Points: {sum(len(seg) for seg in segmentation)}\n"
-        info += f"Simplified Segmentation Points: {sum(len(seg) for seg in simplified_segmentation)}\n"
+        info += f"Original Segmentation Points: {original_points}\n"
+        info += f"Simplified Segmentation Points: {simplified_points}\n"
         info += f"Area: {area:.2f}"
 
         return img, info
@@ -61,7 +72,7 @@ def draw_annotations(index):
 with gr.Blocks() as demo:
     gr.Markdown("# Brain Tumor Image Dataset Viewer")
     gr.Markdown("## Refer to the [dwb2023/brain-tumor-image-dataset-semantic-segmentation](https://huggingface.co/datasets/dwb2023/brain-tumor-image-dataset-semantic-segmentation/viewer/default/test) dataset for more information")
-    gr.Markdown("### Blue: Original Segmentation, Green: Simplified Segmentation (max 20 points)")
+    gr.Markdown("### Red: Bounding Box, Blue: Original Segmentation, Green: Simplified Segmentation (max 20 points)")
     
     with gr.Row():
         with gr.Column(scale=1):