Update app.py

app.py

@@ -8,7 +8,7 @@ from rdp import rdp
 dataset = load_dataset("dwb2023/brain-tumor-image-dataset-semantic-segmentation", split="test")
 
 def simplify_segmentation(segmentation, max_points=20):
-    if not segmentation:
+    if not segmentation or len(segmentation) == 0:
         return []
     epsilon = 1.0
     simplified = rdp(np.array(segmentation), epsilon=epsilon)
@@ -19,38 +19,38 @@ def simplify_segmentation(segmentation, max_points=20):
 
 def draw_annotations(index):
     try:
+        # Fetch the image and annotations from the dataset
         record = dataset[index]
         
+        # Convert image to PIL Image if it's a numpy array
         if isinstance(record['image'], np.ndarray):
             img = Image.fromarray(record['image'])
         else:
             img = record['image']
         
-        img = img.convert("RGB")
+        img = img.convert("RGB")  # Ensure the image is in RGB mode
+        
         draw = ImageDraw.Draw(img)
 
         # Draw bounding box
         bbox = record["bbox"]
         draw.rectangle([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]], outline="red", width=2)
 
-        # Draw segmentation mask
+        # Draw original segmentation mask
         segmentation = record["segmentation"]
         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 if len(seg) > 0]
-            for seg in simplified_segmentation:
-                if len(seg) > 0:
-                    draw.polygon(seg, outline="green", width=2)
-                    simplified_points += len(seg)
+        for seg in segmentation:
+            if seg:  # Check if the segmentation is not empty
+                draw.polygon(seg, outline="blue", width=2)
+                original_points += len(seg)
+                
+                # Simplify and draw simplified segmentation
+                simplified_seg = simplify_segmentation(seg)
+                draw.polygon(simplified_seg, outline="green", width=2)
+                simplified_points += len(simplified_seg)
 
+        # Prepare additional information
         category_id = record["category_id"]
         area = record["area"]
         file_name = record["file_name"]