save raw_image and edited image

app.py

@@ -24,6 +24,7 @@ with block:
     def point_labels_empty():
         return []
     raw_image = gr.Image(type='pil', visible=False)
+    image_edit_trigger = gr.State(True)
     point_coords = gr.State(point_coords_empty)
     point_labels = gr.State(point_labels_empty)
     masks = gr.State([])
@@ -51,7 +52,7 @@ with block:
                 sam_decode_btn = gr.Button('Predict using points!', variant = 'primary')
                 depth_reconstruction_btn = gr.Button('Depth Reconstruction', variant = 'primary')
     # components
-    components = {point_coords, point_labels, raw_image, masks, cutout_idx, input_image,
+    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}
     
@@ -69,7 +70,10 @@ with block:
     input_image.upload(on_input_image_upload, [input_image], [raw_image, point_coords, point_labels], queue=False)
 
     # event - set coords
-    def on_input_image_select(input_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
+    def on_input_image_select(input_image, image_edit_trigger, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
+        if image_edit_trigger:
+            unedited_image = input_image.copy()
+            image_edit_trigger = False
         x, y = evt.index
         color = red if point_label_radio == 0 else blue
         img = np.array(input_image)
@@ -77,8 +81,14 @@ with block:
         img = Image.fromarray(img)
         point_coords.append([x,y])
         point_labels.append(point_label_radio)
-        return img, point_coords, point_labels
-    input_image.select(on_input_image_select, [input_image, point_coords, point_labels, point_label_radio], [input_image, point_coords, point_labels], queue=False)
+        
+        return {raw_image: unedited_image,
+                input_image: img,
+                point_coords: point_coords,
+                point_labels: point_labels,
+                image_edit_trigger: image_edit_trigger}
+    
+    input_image.select(on_input_image_select, [input_image, image_edit_trigger, point_coords, point_labels, point_label_radio], [input_image, raw_image, point_coords, point_labels, image_edit_trigger], queue=False)
 
     def on_click_sam_encode_btn(inputs):
         print("encoding")
@@ -93,6 +103,7 @@ with block:
         image = inputs[raw_image]
         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)