Output mask instead of visual

.gitignore

@@ -0,0 +1 @@
+__pycache__

CutLER

@@ -1 +1 @@
-Subproject commit 077938c626341723050a1971107af552a6ca6697
+Subproject commit 19ca89e1ce39a786e459f64a9d7f6dcfbcdb6d6f

app.py

@@ -6,40 +6,44 @@ import gradio as gr
 import numpy as np
 import PIL.Image as Image
 
-from model import Model, random_color, vis_mask
+from model import Model
+import base64
+from io import BytesIO
 
 model = Model()
 
 
-def run(image_path, threshold, max_num_mask):
+def run(image_path, threshold):
     image = np.asarray(Image.open(image_path).convert('RGB'))
-    masks = model(image_path, threshold, max_num_mask)
-    for mask in masks:
-        image = vis_mask(image, mask, random_color(rgb=True))
-    return image
+    # We copy the image that and fill it with black, to get the dimensions
+    rgb = np.copy(image)
+    rgb.fill(0)
+    masks = model(image_path, threshold, 1)
+
+    mask = masks[0]
+    fg = mask > 0.5
+    rgb[fg] = 255
+    img = Image.fromarray(rgb)
+
+    return img
 
 
 TITLE = 'MaskCut'
 DESCRIPTION = 'This is an unofficial demo for https://github.com/facebookresearch/CutLER.'
 
 paths = sorted(pathlib.Path('CutLER/maskcut/imgs').glob('*.jpg'))
-demo = gr.Interface(
-    fn=run,
-    inputs=[
-        gr.Image(label='Input image', type='filepath'),
-        gr.Slider(label='Threshold used for producing binary graph',
-                  minimum=0,
-                  maximum=1,
-                  value=0.15,
-                  step=0.01),
-        gr.Slider(label='The maximum number of pseudo-masks per image',
-                  minimum=1,
-                  maximum=20,
-                  value=6,
-                  step=1),
-    ],
-    outputs=gr.Image(label='Result', type='numpy'),
-    examples=[[path.as_posix(), 0.15, 6] for path in paths],
-    title=TITLE,
-    description=DESCRIPTION)
-demo.queue().launch()
+demo = gr.Interface(fn=run,
+                    inputs=[
+                        gr.Image(label='Input image', type='filepath'),
+                        gr.Slider(
+                            label='Threshold used for producing binary graph',
+                            minimum=0,
+                            maximum=1,
+                            value=0.15,
+                            step=0.01),
+                    ],
+                    outputs=gr.Image(label='Output image', type="pil"),
+                    examples=[[path.as_posix(), 0.15] for path in paths],
+                    title=TITLE,
+                    description=DESCRIPTION)
+demo.queue().launch(share=True)

model.py

@@ -17,13 +17,6 @@ from maskcut import maskcut
 from third_party.TokenCut.unsupervised_saliency_detection import metric
 
 
-def vis_mask(input, mask, mask_color):
-    fg = mask > 0.5
-    rgb = np.copy(input)
-    rgb[fg] = (rgb[fg] * 0.3 + np.array(mask_color) * 0.7).astype(np.uint8)
-    return Image.fromarray(rgb)
-
-
 class Model:
     def __init__(self):
         self.device = torch.device(