UI change

app.py

@@ -28,7 +28,7 @@ with block:
     point_labels = gr.State(point_labels_empty)
     masks = gr.State([])
     cutout_idx = gr.State(set())
-    everything_masks = gr.State([])
+    pred_masks = gr.State([])
     prompt_masks = gr.State([])
 
     # UI
@@ -36,19 +36,20 @@ with block:
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(label='Input', type='pil', tool=None) # mirror_webcam = False
-                sam_encode_btn = gr.Button('Encode', variant='primary')
+                with gr.Row():
+                    sam_encode_btn = gr.Button('Encode', variant='primary')
+                    sam_sgmt_everything_btn = gr.Button('Segment Everything!', variant = 'primary')
                 #sam_encode_status = gr.Label('Not encoded yet')
         with gr.Row():
-            with gr.Column():
-                prompt_image = gr.Image(label='Segments')
-                with gr.Row():  
-                    point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
-                    text = gr.Textbox(label='Mask Name')
-                    reset_btn = gr.Button('New Mask')
+            prompt_image = gr.Image(label='Segments')
             prompt_lbl_image = gr.AnnotatedImage(label='Segment Labels')
-            everything_image = gr.AnnotatedImage(label='Everything')
+            lbl_image = gr.AnnotatedImage(label='Everything')
+        point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
+        text = gr.Textbox(label='Mask Name')
+        reset_btn = gr.Button('New Mask')
         selected_masks_image = gr.AnnotatedImage(label='Selected Masks')
-        sam_sgmt_everything_btn = gr.Button('Segment Everything!', variant = 'primary')
+       
+        
 
         with gr.Row():
             with gr.Column():
@@ -72,7 +73,7 @@ with block:
         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_prompt_image_select(input_image, prompt_image, point_coords, point_labels, point_label_radio, text, prompt_masks, evt: gr.SelectData):
+    def on_prompt_image_select(input_image, prompt_image, point_coords, point_labels, point_label_radio, text, pred_masks, evt: gr.SelectData):
         x, y = evt.index
         color = red if point_label_radio == 0 else blue
         if prompt_image is None:
@@ -82,32 +83,32 @@ 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))
-        prompt_masks.append((generated_mask, text))
+        pred_masks = [(generated_mask, text)]
         return  [ prompt_image,
                   (input_image, [(generated_mask, "Mask")]),
                   point_coords,
                   point_labels,
-                  prompt_masks ]
+                  pred_masks ]
         
     prompt_image.select(on_prompt_image_select,
-                       [input_image, prompt_image, point_coords, point_labels, point_label_radio, text, prompt_masks],
-                       [prompt_image, prompt_lbl_image, point_coords, point_labels, prompt_masks], queue=False)
+                       [input_image, prompt_image, point_coords, point_labels, point_label_radio, text, pred_masks],
+                       [prompt_image, lbl_image, point_coords, point_labels, pred_masks], queue=False)
     
-    def on_everything_image_select(input_image, everything_masks, masks, text, evt: gr.SelectData):
+    def on_everything_image_select(input_image, pred_masks, masks, text, evt: gr.SelectData):
         i = evt.index
-        mask = everything_masks[i][0]
+        mask = pred_masks[i][0]
         print(mask)
         print(type(mask))
         masks.append((mask, text))
         anno = (input_image, masks) 
         return  [masks, anno]
         
-    everything_image.select(on_everything_image_select,
-                       [input_image, everything_masks, masks, text],
-                       [masks, selected_masks_image], queue=False)
-    prompt_lbl_image.select(on_everything_image_select,
-                       [input_image, prompt_masks, masks, text],
+    lbl_image.select(on_everything_image_select,
+                       [input_image, pred_masks, masks, text],
                        [masks, selected_masks_image], queue=False)
+    #prompt_lbl_image.select(on_everything_image_select,
+    #                   [input_image, prompt_masks, masks, text],
+    #                   [masks, selected_masks_image], queue=False)
 
 
     def on_click_sam_encode_btn(inputs):
@@ -141,7 +142,7 @@ with block:
         print(image)
         print(sam_masks)
         return [(image, sam_masks), sam_masks]
-    sam_sgmt_everything_btn.click(on_sam_sgmt_everything_btn_click, components, [everything_image, everything_masks], queue=False)
+    sam_sgmt_everything_btn.click(on_sam_sgmt_everything_btn_click, components, [lbl_image, pred_masks], queue=False)
 
 
 if __name__ == '__main__':