Update app.py

app.py

@@ -1,8 +1,14 @@
 import gradio as gr
+from PIL import Image as PILImage
+import numpy as np
 from depth_map_generator import generate_depth_map
 from depth_segmentation import segment_image_by_depth
 
 def on_process(image, num_segments):
+    # Convert the NumPy array to a PIL Image
+    if isinstance(image, np.ndarray):
+        image = PILImage.fromarray(image)
+
     # Save the uploaded image
     original_img_path = "original.jpg"
     image.save(original_img_path)
@@ -27,7 +33,7 @@ with gr.Blocks() as demo:
             process_button = gr.Button("Process Image")
         
         with gr.Column():
-            output_gallery = gr.Gallery(label="Segmented Images")
+            output_gallery = gr.Gallery(label="Segmented Images").style(grid=[3, 2])
             download_button = gr.Files(label="Download All", file_count="multiple")
     
     process_button.click(on_process, [image_input, num_segments], [output_gallery, download_button])