Update app.py

app.py

@@ -613,16 +613,17 @@ def infer_seg(media, query):
     path = f"/tmp/{uuid.uuid4().hex}.png"
     iio.imwrite(path, imgs, duration=100, loop=0)
 
-    # masks_value = (media, [(m[0, 0].numpy(), entity_names[i]) for i, m in enumerate(_pred_masks)])
-    # masks_value = (
-    #     media,
-    #     [( _pred_masks[tag_to_mask_idx[tag]][0, 0].numpy(), tag ) for tag in unique_tags]
-    # )
-
+    base_img = Image.open(media).convert("RGB")
     entity_names = [f"Target {i+1}" for i in range(len(unique_tags))]
     masks_value = (
-        media,
-        [(_pred_masks[tag_to_mask_idx[tag]][0, 0].numpy().astype(np.uint8) * 255, entity_names[i]) for i, tag in enumerate(unique_tags)]
+        base_img,
+        [
+            (
+                _pred_masks[tag_to_mask_idx[tag]][0, 0].numpy().astype(bool),
+                entity_names[i]
+            )
+            for i, tag in enumerate(unique_tags)
+        ]
     )
 
     lines = []
@@ -657,7 +658,7 @@ def build_demo():
             msk_1 = gr.AnnotatedImage(label='De-tokenized 2D masks', color_map=color_map, render=False)
             ans_1 = gr.HighlightedText(
                 label='Model Response', color_map=color_map_light, show_inline_category=False, render=False)
-            tag_map_md = gr.Markdown(label="Mask-Token Mapping", value="", visible=False)
+            tag_map_md = gr.Markdown(value="", visible=False)
             with gr.Row():
                 with gr.Column():
                     media_1 = gr.Image(type='filepath')
@@ -685,8 +686,8 @@ def build_demo():
                 
                 with gr.Column():
                     msk_1.render()
-                    tag_map_md
                     ans_1.render()
+                    tag_map_md
 
             ctx_1 = submit_btn_1.click(disable_btns, None, [random_btn_1, reset_btn_1, download_btn_1, submit_btn_1])
             ctx_1 = ctx_1.then(infer_seg, [media_1, query_1], [ans_1, msk_1, download_btn_1, tag_map_md])