Image instead of anno image

app.py

@@ -37,7 +37,7 @@ 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')
-            masks_annotated_image = gr.AnnotatedImage(label='Segments')
+            prompt_image = gr.AnnotatedImage(label='Segments')
             with gr.Column():
                 pcl_figure = gr.Model3D(label="3-D Reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
                 n_samples = gr.Slider(minimum=1e3, maximum=1e6, step=1e3, default=1e3, label='Number of Samples')
@@ -54,7 +54,7 @@ with block:
     # components
     components = {point_coords, point_labels, raw_image, image_edit_trigger, masks, cutout_idx, input_image,
                   point_label_radio, text, reset_btn, sam_sgmt_everything_btn,
-                  sam_decode_btn, depth_reconstruction_btn, masks_annotated_image, n_samples, cube_size}
+                  sam_decode_btn, depth_reconstruction_btn, prompt_image, n_samples, cube_size}
     
     # event - init coords
     def on_reset_btn_click(raw_image):
@@ -86,16 +86,16 @@ with block:
         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))
-        return  [(img, [(generated_mask, text)]),
+        return  [ img,
                   unedited_image,
                   img,
                   point_coords,
                   point_labels,
                   image_edit_trigger]
     
-    masks_annotated_image.select(on_input_image_select,
+    prompt_image.select(on_input_image_select,
                        [raw_image, input_image, image_edit_trigger, point_coords, point_labels, point_label_radio],
-                       [masks_annotated_image, raw_image, input_image, point_coords, point_labels, image_edit_trigger], queue=False)
+                       [prompt_image, raw_image, input_image, point_coords, point_labels, image_edit_trigger], queue=False)
 
     def on_click_sam_encode_btn(inputs):
         print("encoding")
@@ -103,8 +103,8 @@ with block:
         sam.encode(inputs[input_image])
         print("encoding done")
         return {sam_encode_status: 'Image Encoded!',
-                masks_annotated_image: (inputs[input_image], [])}
-    sam_encode_btn.click(on_click_sam_encode_btn, components, [sam_encode_status, masks_annotated_image], queue=False)
+                prompt_image: inputs[input_image]}
+    sam_encode_btn.click(on_click_sam_encode_btn, components, [sam_encode_status, prompt_image], queue=False)
 
     def on_click_sam_dencode_btn(inputs):
         print("inferencing")
@@ -112,8 +112,8 @@ with block:
         generated_mask, _, _ = sam.cond_pred(pts=np.array(inputs[point_coords]), lbls=np.array(inputs[point_labels]))
         inputs[masks].append((generated_mask, inputs[text]))
         print(inputs[masks][0])
-        return {masks_annotated_image: (image, inputs[masks])}
-    sam_decode_btn.click(on_click_sam_dencode_btn, components, [masks_annotated_image, masks, cutout_idx], queue=True)
+        return {prompt_image: (image, inputs[masks])}
+    sam_decode_btn.click(on_click_sam_dencode_btn, components, [prompt_image, masks, cutout_idx], queue=True)
 
     def on_depth_reconstruction_btn_click(inputs):
         print("depth reconstruction")