update app to improve error handling and allow for simultaneous usage

.gitignore

@@ -5,3 +5,4 @@ venv/
 __pycache__/
 .output/
 .data/
+.vscode/

app.py

@@ -26,22 +26,13 @@ ARGS = SimpleNamespace(
 )
 NUM_SAMPLES = 10
 
-outputs = []
+outputs = {}
 
 
 def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_samples: int) -> list[Any]:
     global outputs
 
-    def find_gifs(path: str) -> list[str]:
-        """Scrape folders for all generated gif files."""
-        for file in os.listdir(path):
-            sub_path = os.path.join(path, file)
-            if os.path.isdir(sub_path):
-                for image_file in os.listdir(sub_path):
-                    if re.match(r".*\.gif$", image_file):
-                        yield os.path.join(sub_path, image_file)
-
-    def find_images(path: str) -> list[str]:
+    def find_images(path: str) -> dict[str, list[str]]:
         """Scrape folders for all generated gif files."""
         images = {}
         for file in os.listdir(path):
@@ -62,6 +53,14 @@ def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_sample
         else:
             os.remove(full_path)
 
+    if not rgb_image:
+        gr.Error("You must provide an RGB image before running the model.")
+        return [None] * 5
+
+    if not depth_image:
+        gr.Error("You must provide a depth image before running the model.")
+        return [None] * 5
+
     cfg = setup_cfg(ARGS)
 
     engine.launch(
@@ -80,35 +79,42 @@ def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_sample
 
     # process output
     # TODO: may want to select these in decreasing order of score
+    outputs[rgb_image] = []
     image_files = find_images(ARGS.output)
-    outputs = []
     for count, part in enumerate(image_files):
         if count < MAX_PARTS:
-            outputs.append([Image.open(im) for im in image_files[part]])
+            outputs[rgb_image].append([Image.open(im) for im in image_files[part]])
 
     return [
-        *[gr.update(value=out[0], visible=True) for out in outputs],
+        *[gr.update(value=out[0], visible=True) for out in outputs[rgb_image]],
         *[gr.update(visible=False) for _ in range(MAX_PARTS - len(outputs))],
     ]
 
 
 def get_trigger(idx: int, fps: int = 40, oscillate: bool = True):
-    def iter_images(*args, **kwargs):
-        if idx < len(outputs):
-            for im in outputs[idx]:
+    def iter_images(rgb_image: str):
+        if not rgb_image or rgb_image not in outputs:
+            gr.Warning("You must upload an image and run the model before you can view the output.")
+
+        elif idx < len(outputs[rgb_image]):
+            for im in outputs[rgb_image][idx]:
                 time.sleep(1.0 / fps)
                 yield im
             if oscillate:
-                for im in reversed(outputs[idx]):
+                for im in reversed(outputs[rgb_image][idx]):
                     time.sleep(1.0 / fps)
                     yield im
 
         else:
-            raise ValueError("Could not find any images to load into this module.")
+            gr.Error("Could not find any images to load into this module.")
 
     return iter_images
 
 
+def clear_outputs():
+    return [gr.update(value=None, visible=(idx == 0)) for idx in range(MAX_PARTS)]
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -176,12 +182,20 @@ with gr.Blocks() as demo:
     )
 
     submit_btn = gr.Button("Run model")
+    explanation = gr.Markdown(value="# Output\nClick on an image to see an animation of the part motion.")
 
     # TODO: do we want to set a maximum limit on how many parts we render? We could also show the number of components
     # identified.
-    images = [gr.Image(type="pil", label=f"Part {idx + 1}", visible=False) for idx in range(MAX_PARTS)]
+    images = [
+        gr.Image(type="pil", label=f"Part {idx + 1}", show_download_button=False, visible=(idx == 0))
+        for idx in range(MAX_PARTS)
+    ]
     for idx, image_comp in enumerate(images):
-        image_comp.select(get_trigger(idx), inputs=[], outputs=image_comp, api_name=False)
+        image_comp.select(get_trigger(idx), inputs=rgb_image, outputs=image_comp, api_name=False)
+
+    # if user changes input, clear output images
+    rgb_image.change(clear_outputs, inputs=[], outputs=images, api_name=False)
+    depth_image.change(clear_outputs, inputs=[], outputs=images, api_name=False)
 
     submit_btn.click(
         fn=predict, inputs=[rgb_image, depth_image, intrinsics, num_samples], outputs=images, api_name=False