V2

app.py

@@ -89,29 +89,39 @@ probs = np.zeros((11, 128, 128), dtype=np.float32)
 
 with gr.Blocks() as demo:
     
+    # work in progress
+    # with gr.Row():
+    #     gr_gallery = gr.Gallery(value=[f"images/img{i}.png" for i in range(1, 9)])
+    #     gr_gallery = gr_gallery.style(columns=[3], rows=[3], object_fit="contain", height="auto")
+
     with gr.Row():
 
-      with gr.Column():
         gr_choose_image = gr.Dropdown(
             [f"img{i}" for i in range(1, 9)], label="CLEVR Image", info="Start by a picking an image from the CLEVR dataset."
         )
-        gr_image_1 = gr.Image(type="numpy")
-        gr_image_2 = gr.Image(type="numpy")
+
+    with gr.Row():
 
       with gr.Column():
-        gr_slot_slider = gr.Slider(1, 11, value=1, step=1, label="Slot")
+
+        gr_image_1 = gr.Image(type="numpy", source="canvas", label="Decoding")
+        gr_image_2 = gr.Image(type="numpy", source="canvas", label="Segmentation")
+
+      with gr.Column():
+        gr_slot_slider = gr.Slider(1, 11, value=1, step=1, label="Slot Index",
+                                   info="Change slot index too see the segmentation mask, position and scale of each slot.")
 
         gr_y_slider = gr.Slider(-1, 1, value=0, step=0.01, label="x")
         gr_x_slider = gr.Slider(-1, 1, value=0, step=0.01, label="y")
         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")
+        gr_button = gr.Button("Render", info="Render a new image with altered positions and scales.")
 
     def update_image_and_segmentation(name, idx):
       idx = idx - 1
 
-      img_input, img = load_image(name)
+      img_input, _ = load_image(name)
       out = model.apply(
         {"params": state.params, **state.variables},
         video=img_input[None, None],
@@ -119,12 +129,15 @@ with gr.Blocks() as demo:
         train=False)
       
       probs[:] = nn.softmax(out["outputs"]["segmentation_logits"][0, 0, :, :, :, 0], axis=0)
+      img = np.array(out["outputs"]["video"][0, 0])
+      img = np.clip(img, 0, 1)
+
       slots_ = out["states"]
       slots[:] = slots_[0, 0, :, :-4]
       pos[:] = slots_[0, 0, :, -4: -2]
       scale[:] = slots_[0, 0, :, -2:]
 
-      return img, (probs[idx] * 255).astype(np.uint8), float(pos[idx, 0]), \
+      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])
 
     gr_choose_image.change(