Allow for multiple objects

app.py

@@ -13,20 +13,19 @@ from io import BytesIO
 model = Model()
 
 
-def run(image_path, threshold):
+def run(image_path, threshold, num_objects):
     image = Image.open(image_path).convert('RGB')
     canvas = Image.new('RGB', image.size, (0, 0, 0))
 
     # We copy the image that and fill it with black, to get the dimensions
     rgb = np.array(canvas)
-    masks = model(image_path, threshold, 1)
+    masks = model(image_path, threshold, num_objects)
 
-    mask = masks[0]
-    fg = mask > 0.5
-    rgb[fg] = 255
-    img = Image.fromarray(rgb)
+    for mask in masks:  
+      fg = mask > 0.5
+      rgb[fg] = 255
 
-    return img
+    return Image.fromarray(rgb)
 
 
 TITLE = 'MaskCut'
@@ -42,9 +41,15 @@ demo = gr.Interface(fn=run,
                             maximum=1,
                             value=0.15,
                             step=0.01),
+                        gr.Slider(
+                            label='Number of objects in the image',
+                            minimum=1,
+                            maximum=10,
+                            value=1,
+                            step=1)
                     ],
                     outputs=gr.Image(label='Output image', type="pil"),
-                    examples=[[path.as_posix(), 0.15] for path in paths],
+                    examples=[[path.as_posix(), 0.15, 2] for path in paths],
                     title=TITLE,
                     description=DESCRIPTION)
 demo.queue(concurrency_count=3, api_open=True)