Update app.py

app.py

@@ -5,12 +5,22 @@ import numpy as np
 
 # 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 draw_annotations(index):
+# Function to filter dataset based on category ID
+def filter_dataset_by_category(category_id):
+    filtered_indices = [i for i, record in enumerate(dataset) if record["category_id"] == category_id]
+    return filtered_indices
+
+# Function to draw annotations
+def draw_annotations(index, category_id):
+    filtered_indices = filter_dataset_by_category(category_id)
+    
+    if index >= len(filtered_indices):
+        index = 0
+    
     try:
         # Fetch the image and annotations from the dataset
-        record = dataset[index]
+        record = dataset[filtered_indices[index]]
         
         # Convert image to PIL Image if it's a numpy array
         if isinstance(record['image'], np.ndarray):
@@ -43,10 +53,10 @@ def draw_annotations(index):
         info += f"Segmentation: {segmentation}\n"
         info += f"Area: {area:.2f}"
 
-        return img, info
+        return img, info, len(filtered_indices) - 1
     except Exception as e:
         print(f"Error processing image at index {index}: {e}")
-        return Image.new('RGB', (300, 300), color='gray'), f"Error loading image information: {str(e)}"
+        return Image.new('RGB', (300, 300), color='gray'), f"Error loading image information: {str(e)}", len(filtered_indices) - 1
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -57,13 +67,20 @@ with gr.Blocks() as demo:
         with gr.Column(scale=1):
             image_output = gr.Image(label="Annotated Image")
         with gr.Column(scale=1):
-            image_index = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, value=0, label="Image ID Slider")
+            category_id_dropdown = gr.Dropdown(choices=[1, 2], value=1, label="Category ID")
+            image_index = gr.Slider(minimum=0, maximum=0, step=1, value=0, label="Image ID Slider")
             info_output = gr.Textbox(label="Image Information", lines=10)
     
-    # Update image and info when slider changes
-    image_index.change(draw_annotations, inputs=image_index, outputs=[image_output, info_output])
+    def update_slider(category_id):
+        _, _, max_index = draw_annotations(0, category_id)
+        return gr.Slider.update(maximum=max_index)
+
+    # Update image and info when slider or category changes
+    category_id_dropdown.change(update_slider, inputs=category_id_dropdown, outputs=image_index)
+    category_id_dropdown.change(draw_annotations, inputs=[image_index, category_id_dropdown], outputs=[image_output, info_output, image_index])
+    image_index.change(draw_annotations, inputs=[image_index, category_id_dropdown], outputs=[image_output, info_output, image_index])
 
     # Display initial image and info
-    demo.load(draw_annotations, inputs=image_index, outputs=[image_output, info_output])
+    demo.load(draw_annotations, inputs=[image_index, category_id_dropdown], outputs=[image_output, info_output, image_index])
 
-demo.launch(debug=True)
+demo.launch(debug=True)