fix

app.py

@@ -39,12 +39,33 @@ model = model.eval()
 
 import gradio as gr
 
+colors = [
+    (0, 255, 0),
+    (0, 0, 255),
+    (255, 255, 0),
+    (255, 0, 255),
+    (0, 255, 255),
+    (114, 128, 250),
+    (0, 165, 255),
+    (0, 128, 0),
+    (144, 238, 144),
+    (238, 238, 175),
+    (255, 191, 0),
+    (0, 128, 0),
+    (226, 43, 138),
+    (255, 0, 255),
+    (0, 215, 255),
+    (255, 0, 0),
+]
+
+color_map = {
+    f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}" for color_id, color in enumerate(colors)
+}
+
 def run_clip_dinoiser(input_image, text_prompts):
     image = input_image.convert("RGB")
     text_prompts = text_prompts.split(",")
-    palette = [
-        (random.randint(0, 256), random.randint(0, 256), random.randint(0, 256)) for _ in range(len(text_prompts))
-    ]
+    palette = colors[:len(text_prompts)]
 
     model.clip_backbone.decode_head.update_vocab(text_prompts)
     model.to(device)
@@ -78,7 +99,14 @@ def run_clip_dinoiser(input_image, text_prompts):
     palette_array = np.array(itemgetter(*classes)(palette)).reshape(1, -1, 3)
     alpha=0.5
     blend = (alpha)*np.array(image)/255. + (1-alpha) * mask/255.
-    return palette_array, blend, mask
+
+    h_text = list()
+    for idx, text in enumerate(text_prompts):
+        for alphabet in text:
+            h_text.append((alphabet, color_map[str(idx)]))
+    return blend, mask, h_text
+
+
 
 
 if __name__ == "__main__":
@@ -94,9 +122,6 @@ if __name__ == "__main__":
                 run_button = gr.Button(value="Run")
 
             with gr.Column():
-                palette_array = gr.Image(
-                    type="numpy",
-                )
                 with gr.Row():
                     overlay_mask = gr.Image(
                         type="numpy",
@@ -104,18 +129,24 @@ if __name__ == "__main__":
                     only_mask = gr.Image(
                         type="numpy",
                     )
+                h_text = gr.HighlightedText(
+                    label="text",
+                    combine_adjacent=False,
+                    show_legend=False,
+                    color_map=color_map
+                )
 
         run_button.click(
             fn=run_clip_dinoiser,
             inputs=[input_image, text_prompts,],
-            outputs=[palette_array, overlay_mask, only_mask]
+            outputs=[overlay_mask, only_mask, h_text]
         )
         gr.Examples(
-          [["vintage_bike.jpeg", "background, vintage bike, leather bag"]],
-          inputs = [input_image, text_prompts,],
-          outputs = [palette_array, overlay_mask, only_mask],
-          fn=run_clip_dinoiser,
-          cache_examples=True,
-          label='Try this example input!'
+            [["vintage_bike.jpeg", "background, vintage bike, leather bag"]],
+            inputs = [input_image, text_prompts,],
+            outputs = [overlay_mask, only_mask, h_text],
+            fn=run_clip_dinoiser,
+            cache_examples=True,
+            label='Try this example input!'
       )
     block.launch(share=False, show_api=False, show_error=True)