Update app.py

app.py

@@ -548,6 +548,7 @@ def infer_seg(media, query):
             answer,
             gr.update(value=None, visible=False),
             gr.update(value=None, interactive=False, visible=False),
+            gr.update(value="", visible=False),
         )
 
     batch_size = len(quant_ids) // CODEBOOK_DEPTH
@@ -597,34 +598,18 @@ def infer_seg(media, query):
 
     answer = dict(text=output_text, entities=entities)
 
-    # entities = []
-    # unique_tags = list(set(tags))
-    # entity_names = []
-    # for i, tag in enumerate(unique_tags):
-    #     for m in re.finditer(re.escape(tag), output_text):
-    #         entities.append(dict(entity=f'Target {i + 1}', start=m.start(), end=m.end()))
-    #         entity_names.append(f'Target {i + 1}')
-    
-    # answer = dict(text=output_text, entities=entities)
-
     frames = torch.from_numpy(np.array(image)).unsqueeze(0)
     imgs = draw_mask(frames, _pred_masks, colors=colors)
 
     path = f"/tmp/{uuid.uuid4().hex}.png"
     iio.imwrite(path, imgs, duration=100, loop=0)
 
-    base_img = Image.open(media).convert("RGB")
+    mask_items = []
     entity_names = [f"Target {i+1}" for i in range(len(unique_tags))]
-    masks_value = (
-        base_img,
-        [
-            (
-                _pred_masks[tag_to_mask_idx[tag]][0, 0].numpy().astype(bool),
-                entity_names[i]
-            )
-            for i, tag in enumerate(unique_tags)
-        ]
-    )
+    for i, tag in enumerate(unique_tags):
+        m = _pred_masks[tag_to_mask_idx[tag]][0, 0].numpy().astype(np.uint8) * 255  # (H,W), 0/255
+        mask_items.append((m, entity_names[i]))
+    masks_value = (media, mask_items)
 
     lines = []
     for i, tag in enumerate(unique_tags):
@@ -658,7 +643,6 @@ 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(value="", visible=False)
             with gr.Row():
                 with gr.Column():
                     media_1 = gr.Image(type='filepath')
@@ -677,7 +661,7 @@ def build_demo():
                     with gr.Row():
                         random_btn_1 = gr.Button(value='🔮 Random', visible=False)
 
-                        reset_btn_1 = gr.ClearButton([media_1, query_1, msk_1, ans_1, tag_map_md], value='🗑️ Reset')
+                        reset_btn_1 = gr.ClearButton([media_1, query_1, msk_1, ans_1], value='🗑️ Reset')
                         reset_btn_1.click(reset_seg, None, [sample_frames_1, download_btn_1])
 
                         download_btn_1.render()
@@ -687,7 +671,7 @@ def build_demo():
                 with gr.Column():
                     msk_1.render()
                     ans_1.render()
-                    tag_map_md
+                    tag_map_md = gr.Markdown(value="", visible=False)
 
             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])