Select image from everything annotation

app.py

@@ -48,6 +48,7 @@ with block:
                     n_samples = gr.Slider(minimum=1e3, maximum=1e6, step=1e3, default=1e3, label='Number of Samples')
                     cube_size = gr.Slider(minimum=0.000001, maximum=0.001, step=0.000001, default=0.00001, label='Cube size')
         with gr.Row():
+            selected_masks_image = gr.AnnotatedImage(label='Selected Masks')
             with gr.Column(scale=1):
                 with gr.Row():
                     point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
@@ -59,22 +60,13 @@ with block:
     # components
     components = {point_coords, point_labels, image_edit_trigger, masks, cutout_idx, input_image,
                   point_label_radio, text, reset_btn, sam_sgmt_everything_btn,
-                  sam_decode_btn, depth_reconstruction_btn, prompt_image, prompt_lbl_image, n_samples, cube_size}
+                  sam_decode_btn, depth_reconstruction_btn, prompt_image, prompt_lbl_image, n_samples, cube_size, selected_masks_image}
     
     # event - init coords
     def on_reset_btn_click(input_image):
         return input_image, point_coords_empty(), point_labels_empty(), None, []
     reset_btn.click(on_reset_btn_click, [input_image], [input_image, point_coords, point_labels], queue=False)
 
-    #def on_input_image_upload(input_image):
-    #    print("encoding")
-    #    # encode image on upload
-    ##    sam.encode(input_image)
-    #    print("encoding done")
-    #    return input_image, point_coords_empty(), point_labels_empty(), None
-    #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, prompt_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
         x, y = evt.index
         color = red if point_label_radio == 0 else blue
@@ -94,6 +86,18 @@ with block:
     prompt_image.select(on_prompt_image_select,
                        [input_image, prompt_image, point_coords, point_labels, point_label_radio],
                        [prompt_image, prompt_lbl_image, point_coords, point_labels], queue=False)
+    
+    def on_everything_image_select(input_image ,everything_image, masks, text, evt: gr.SelectData):
+        i = evt.index
+        mask = everything_image[1][i]
+        masks.append((mask, text))
+        anno = (input_image, masks) 
+        return  [ masks, anno]
+        
+    prompt_image.select(on_everything_image_select,
+                       [input_image, everything_image, masks, text],
+                       [masks, selected_masks_image], queue=False)
+
 
     def on_click_sam_encode_btn(inputs):
         print("encoding")