Update app.py

app.py

@@ -7,6 +7,11 @@ import numpy as np
 dataset = load_dataset("dwb2023/brain-tumor-image-dataset-semantic-segmentation", split="test")
 # print(f"Dataset loaded successfully. Number of images: {len(dataset)}")
 
+CATEGORY_COLORS = {
+    1: "red",
+    2: "green"
+}
+
 def draw_annotations(index):
     try:
         # Fetch the image and annotations from the dataset
@@ -22,9 +27,11 @@ def draw_annotations(index):
         
         draw = ImageDraw.Draw(img)
 
-        # Draw bounding box
+        # Draw bounding box with color based on category
         bbox = record["bbox"]
-        draw.rectangle([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]], outline="red", width=2)
+        category_id = record["category_id"]
+        box_color = CATEGORY_COLORS.get(category_id, "blue")  # Default to blue if category not in mapping
+        draw.rectangle([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]], outline=box_color, width=2)
 
         # Draw segmentation mask
         segmentation = record["segmentation"]
@@ -32,13 +39,13 @@ def draw_annotations(index):
             draw.polygon(seg, outline="blue", width=2)
 
         # Prepare additional information
-        category_id = record["category_id"]
         area = record["area"]
         file_name = record["file_name"]
 
         info = f"File Name: {file_name}\n"
         info += f"Image ID: {record['id']}\n"
         info += f"Category ID: {category_id}\n"
+        info += f"Bounding Box Color: {box_color}\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"Area: {area:.2f}"