Update app.py

app.py

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