stack masks

app.py

@@ -47,6 +47,7 @@ with gr.Blocks() as demo:
     gr.Markdown("# WildSynth: Synthetic Wildlife Data Generation")
     gr.Markdown(
         """
+        ## Work in Progress
         ### About
         We have trained a JAX ControlNet model for semantic segmentation on Wildlife Animal Images.
         
@@ -73,13 +74,14 @@ with gr.Blocks() as demo:
 
     def generate_mask(image):
         outputs = generator(image, points_per_batch=256)
-        
+        mask_images = []
         for mask in outputs["masks"]:
             color = np.concatenate([np.random.random(3), np.array([1.0])], axis=0)
             h, w = mask.shape[-2:]
             mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            mask_images.append(mask_image)
         
-        return mask_image
+        return np.stack(mask_images)
 
         # predictor.set_image(image)
         # input_point = np.array([120, 21])