Added a reset button.

app.py

@@ -77,10 +77,15 @@ class DecoderWrapper(nn.Module):
 decoder_model = DecoderWrapper(decoder=model.decoder)
 
 with gr.Blocks() as demo:
-    
+
     local_slots = gr.State(np.zeros((11, 64), dtype=np.float32))
+
+    local_orig_pos = gr.State(np.zeros((11, 2), dtype=np.float32))
+    local_orig_scale = gr.State(np.zeros((11, 2), dtype=np.float32))
+
     local_pos = gr.State(np.zeros((11, 2), dtype=np.float32))
     local_scale = gr.State(np.zeros((11, 2), dtype=np.float32))
+
     local_probs = gr.State(np.zeros((11, 128, 128), dtype=np.float32))
 
     with gr.Row():
@@ -91,7 +96,7 @@ with gr.Blocks() as demo:
 
     with gr.Row():
 
-      with gr.Column():
+      with gr.Column(min_width=600):
 
         with gr.Row():
           with gr.Column():
@@ -108,7 +113,11 @@ with gr.Blocks() as demo:
         gr_sy_slider = gr.Slider(0.01, 1, value=0.1, step=0.01, label="width")
         gr_sx_slider = gr.Slider(0.01, 1, value=0.1, step=0.01, label="height")
 
-        gr_button = gr.Button("Render", info="Render a new image with altered positions and scales.")
+        with gr.Row():
+          with gr.Column():
+            gr_button_render = gr.Button("Render", variant="primary", info="Render a new image with altered positions and scales.")
+          with gr.Column():
+            gr_button_reset = gr.Button("Reset", info="Reset slot statistics.")
 
     def update_image_and_segmentation(name, idx):
       idx = idx - 1
@@ -130,13 +139,13 @@ with gr.Blocks() as demo:
       scale = slots_[0, 0, :, -2:]
 
       return (img * 255).astype(np.uint8), (probs[idx] * 255).astype(np.uint8), float(pos[idx, 0]), \
-             float(pos[idx, 1]), float(scale[idx, 0]), float(scale[idx, 1]), probs, slots, pos, scale
+             float(pos[idx, 1]), float(scale[idx, 0]), float(scale[idx, 1]), probs, slots, pos, scale, pos, scale
 
     gr_choose_image.change(
        fn=update_image_and_segmentation,
        inputs=[gr_choose_image, gr_slot_slider],
        outputs=[gr_image_1, gr_image_2, gr_x_slider, gr_y_slider, gr_sx_slider, gr_sy_slider,
-                local_probs, local_slots, local_pos, local_scale]
+                local_probs, local_slots, local_pos, local_scale, local_orig_pos, local_orig_scale]
     )
 
     def update_sliders(idx, local_probs, local_pos, local_scale):
@@ -151,18 +160,22 @@ with gr.Blocks() as demo:
     )
 
     def update_pos_x(idx, val, local_pos):
+       local_pos = np.copy(local_pos)
        local_pos[idx - 1, 0] = val
        return local_pos
 
     def update_pos_y(idx, val, local_pos):
+       local_pos = np.copy(local_pos)
        local_pos[idx - 1, 1] = val
        return local_pos
 
     def update_scale_x(idx, val, local_scale):
+       local_scale = np.copy(local_scale)
        local_scale[idx - 1, 0] = val
        return local_scale
 
     def update_scale_y(idx, val, local_scale):
+       local_scale = np.copy(local_scale)
        local_scale[idx - 1, 1] = val
        return local_scale
 
@@ -204,10 +217,21 @@ with gr.Blocks() as demo:
       image = np.clip(image, 0, 1)
       return (image * 255).astype(np.uint8), (probs[idx] * 255).astype(np.uint8), probs
 
-    gr_button.click(
+    gr_button_render.click(
         fn=render,
         inputs=[gr_slot_slider, local_slots, local_pos, local_scale],
         outputs=[gr_image_1, gr_image_2, local_probs]
     )
 
+    def reset(idx, local_orig_pos, local_orig_scale):
+       idx = idx - 1
+       return np.copy(local_orig_pos), np.copy(local_orig_scale), float(local_orig_pos[idx, 0]), \
+              float(local_orig_pos[idx, 1]), float(local_orig_scale[idx, 0]), float(local_orig_scale[idx, 1])
+
+    gr_button_reset.click(
+        fn=reset,
+        inputs=[gr_slot_slider, local_orig_pos, local_orig_scale],
+        outputs=[local_pos, local_scale, gr_x_slider, gr_y_slider, gr_sx_slider, gr_sy_slider]
+    )
+
 demo.launch()