Update app.py

app.py

@@ -2,36 +2,44 @@ import gradio as gr
 from datasets import load_dataset
 from PIL import Image, ImageDraw
 import numpy as np
+from rdp import rdp
 
 # Load the dataset
 dataset = load_dataset("dwb2023/brain-tumor-image-dataset-semantic-segmentation", split="test")
-# print(f"Dataset loaded successfully. Number of images: {len(dataset)}")
+
+def simplify_segmentation(segmentation, max_points=20):
+    simplified = rdp(np.array(segmentation), epsilon=1.0)
+    while len(simplified) > max_points:
+        epsilon *= 1.5
+        simplified = rdp(np.array(segmentation), epsilon=epsilon)
+    return simplified.tolist()
 
 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")  # Ensure the image is in RGB mode
-        
+        img = img.convert("RGB")
         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"]
         for seg in segmentation:
             draw.polygon(seg, outline="blue", width=2)
 
-        # Prepare additional information
+        # 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)
+
         category_id = record["category_id"]
         area = record["area"]
         file_name = record["file_name"]
@@ -40,7 +48,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"Segmentation: {segmentation}\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"Area: {area:.2f}"
 
         return img, info
@@ -52,6 +61,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)")
     
     with gr.Row():
         with gr.Column(scale=1):