Update app.py

app.py

@@ -13,7 +13,7 @@ def process_image(image, num_segments):
     # Perform depth segmentation
     segmented_images = segment_image_by_depth(original_img_path, depth_map_path, num_segments)
     
-    return segmented_images
+    return segmented_images, segmented_images
 
 # Gradio Interface
 with gr.Blocks() as demo:
@@ -26,13 +26,9 @@ with gr.Blocks() as demo:
             process_button = gr.Button("Process Image")
         
         with gr.Column():
-            output_gallery = gr.Gallery(label="Segmented Images").style(grid=[3, 2])
+            output_gallery = gr.Gallery(label="Segmented Images")
             download_button = gr.Files(label="Download All", file_count="multiple")
     
-    def on_process(image, num_segments):
-        segmented_images = process_image(image, num_segments)
-        return segmented_images, segmented_images
-
     process_button.click(on_process, [image_input, num_segments], [output_gallery, download_button])
 
 demo.launch()