Ui update

app.py

@@ -36,7 +36,16 @@ with block:
                 input_image = gr.Image(label='Input', type='pil', tool=None) # mirror_webcam = False
                 sam_encode_btn = gr.Button('Encode', variant='primary')
                 sam_encode_status = gr.Label('Not encoded yet')
-            prompt_image = gr.Image(label='Segments')
+        with gr.Row():
+            with gr.Tab("Select with points"):
+                with gr.Column():
+                    prompt_image = gr.Image(label='Segments')
+                    prompt_lbl_image = gr.AnnotatedImage(label='Segment Labels')
+                    
+                
+            with gr.Tab("Select from segmented map"):
+                everything_image = gr.AnnotatedImage(label='Everything')
+            
         with gr.Row():
             with gr.Column():
                 pcl_figure = gr.Model3D(label="3-D Reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
@@ -71,24 +80,25 @@ with block:
     input_image.upload(on_input_image_upload, [input_image], [input_image, point_coords, point_labels], queue=False)
 
     # event - set coords
-    def on_prompt_image_select(input_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
+    def on_prompt_image_select(input_image, prompt, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
         x, y = evt.index
         color = red if point_label_radio == 0 else blue
-        img = input_image.copy()
-        img = np.array(img)
-        cv2.circle(img, (x, y), 5, color, -1)
+        if prompt is None:
+            prompt = np.array(input_image.copy())
+
+        cv2.circle(prompt, (x, y), 5, color, -1)
         point_coords.append([x,y])
         point_labels.append(point_label_radio)
         generated_mask, _, _ = sam.cond_pred(pts=np.array(point_coords), lbls=np.array(point_labels))
-        img[generated_mask] = (255.0, 0.0, 0.0)
-        img = Image.fromarray(img)
-        return  [ img,
+
+        return  [ prompt,
+                  (input_image, [(generated_mask, "Mask")]),
                   point_coords,
                   point_labels ]
-    
+        
     prompt_image.select(on_prompt_image_select,
                        [input_image, point_coords, point_labels, point_label_radio],
-                       [prompt_image, point_coords, point_labels], queue=False)
+                       [prompt_image, prompt_lbl_image, point_coords, point_labels], queue=False)
 
     def on_click_sam_encode_btn(inputs):
         print("encoding")