finish documentation of repo

README.md

@@ -9,7 +9,13 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# OPDMulti: Openable Part Detection for Multiple Objects
+[Xiaohao Sun*](https://sun-xh.github.io/), [Hanxiao Jiang*](https://jianghanxiao.github.io/), [Manolis Savva](https://msavva.github.io/), [Angel Xuan Chang](http://angelxuanchang.github.io/)
+
+This repository is intended as a deployment of a demo for the [OPDMulti](https://github.com/3dlg-hcvc/OPDMulti) project.
+Please refer there for more information about the proect and implementation.
+
+[arXiv](https://arxiv.org/abs/2303.14087)  [Website](https://3dlg-hcvc.github.io/OPDMulti/)
 
 ## Installation
 
@@ -90,3 +96,40 @@ gradio app.py
 You can view the app on the specified port (usually 7860). To run over an ssh connection, setup port forwarding using
 `-L 7860:localhost:7860` when you create your ssh connection. Note that you will need to install Open3D in headless 
 rendering for this to work, as described above.
+
+## Citation
+If you find this code useful, please consider citing:
+```bibtex
+@article{sun2023opdmulti,
+  title={OPDMulti: Openable Part Detection for Multiple Objects},
+  author={Sun, Xiaohao and Jiang, Hanxiao and Savva, Manolis and Chang, Angel Xuan},
+  journal={arXiv preprint arXiv:2303.14087},
+  year={2023}
+}
+
+@article{mao2022multiscan,
+  title={MultiScan: Scalable RGBD scanning for 3D environments with articulated objects},
+  author={Mao, Yongsen and Zhang, Yiming and Jiang, Hanxiao and Chang, Angel and Savva, Manolis},
+  journal={Advances in Neural Information Processing Systems},
+  volume={35},
+  pages={9058--9071},
+  year={2022}
+}
+
+@inproceedings{jiang2022opd,
+  title={OPD: Single-view 3D openable part detection},
+  author={Jiang, Hanxiao and Mao, Yongsen and Savva, Manolis and Chang, Angel X},
+  booktitle={Computer Vision--ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23--27, 2022, Proceedings, Part XXXIX},
+  pages={410--426},
+  year={2022},
+  organization={Springer}
+}
+
+@inproceedings{cheng2022masked,
+  title={Masked-attention mask transformer for universal image segmentation},
+  author={Cheng, Bowen and Misra, Ishan and Schwing, Alexander G and Kirillov, Alexander and Girdhar, Rohit},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={1290--1299},
+  year={2022}
+}
+```

app.py

@@ -3,11 +3,12 @@ import re
 import shutil
 import time
 from types import SimpleNamespace
-from typing import Any
+from typing import Any, Callable, Generator, Optional
 
 import gradio as gr
 import numpy as np
 from detectron2 import engine
+from natsort import natsorted
 from PIL import Image
 
 from inference import main, setup_cfg
@@ -16,7 +17,7 @@ from inference import main, setup_cfg
 NUM_PROCESSES = 1
 CROP = False
 SCORE_THRESHOLD = 0.8
-MAX_PARTS = 5
+MAX_PARTS = 5  # TODO: we can replace this by having a slider and a single image visualization component rather than multiple components
 ARGS = SimpleNamespace(
     config_file="configs/coco/instance-segmentation/swin/opd_v1_real.yaml",
     model=".data/models/motion_state_pred_opdformerp_rgb.pth",
@@ -26,25 +27,49 @@ ARGS = SimpleNamespace(
 )
 NUM_SAMPLES = 10
 
-outputs = {}
-
-
-def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_samples: int) -> list[Any]:
+# this variable holds the current state of results, as the user will need to be able to "reload" the results in order
+# to visualize the demo again. The output images are cached by the temporary path of the image, meaning that multiple
+# users should be able to simultaneously run the demo. Gradio should be able to handle the case where multiple distinct
+# images are uploaded with the same name, as I believe the caching of temp path is based on base64 encoding, not the
+# filename itself.
+# TODO: right now there is no gc system for outputs, which means if there is enough traffic per unit time such that the
+# outputs are all generated on the same system instantiation of the code, the RAM could max out, acknowledging also that
+# this is not designed to run on GPU and so the model and all will also need to be stored in CPU memory. Solutions could
+# include
+#  1. a caching design to remove old results periodically, especially if the image is reset;
+#  2. caching results on disk rather than in memory, since the cap is higher; or
+#  3. figuring out some way to cache results in browser instead of in the backend (couldn't figure out a way to do this
+#     earlier.
+outputs: dict[str, list[list[Image.Image]]] = {}
+
+
+def predict(rgb_image: str, depth_image: str, intrinsic: np.ndarray, num_samples: int) -> list[Any]:
+    """
+    Run model on input image and generate output visualizations.
+
+    :param rgb_image: local path to RGB image file, used for model prediction and visualization
+    :param depth_image: local path to depth image file, used for visualization
+    :param intrinsic: array of dimension (3, 3) representing the intrinsic matrix of the camera
+    :param num_samples: number of visualization states to generate.
+    :return: list of updates to make to image components to visualize first image of visualization sequence, or
+    otherwise to hide an image component from visualization.
+    """
     global outputs
 
     def find_images(path: str) -> dict[str, list[str]]:
-        """Scrape folders for all generated gif files."""
+        """Scrape folders for all generated image files."""
         images = {}
         for file in os.listdir(path):
             sub_path = os.path.join(path, file)
             if os.path.isdir(sub_path):
                 images[file] = []
-                for image_file in sorted(os.listdir(sub_path)):
+                for image_file in natsorted(os.listdir(sub_path)):
                     if re.match(r".*\.png$", image_file):
                         images[file].append(os.path.join(sub_path, image_file))
         return images
 
     # clear old predictions
+    # TODO: might be a better place for this than at the beginning of every invocation
     os.makedirs(ARGS.output, exist_ok=True)
     for path in os.listdir(ARGS.output):
         full_path = os.path.join(ARGS.output, path)
@@ -61,8 +86,8 @@ def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_sample
         gr.Error("You must provide a depth image before running the model.")
         return [None] * 5
 
+    # run model
     cfg = setup_cfg(ARGS)
-
     engine.launch(
         main,
         NUM_PROCESSES,
@@ -70,7 +95,7 @@ def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_sample
             cfg,
             rgb_image,
             depth_image,
-            intrinsics,
+            intrinsic,
             num_samples,
             CROP,
             SCORE_THRESHOLD,
@@ -82,7 +107,7 @@ def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_sample
     outputs[rgb_image] = []
     image_files = find_images(ARGS.output)
     for count, part in enumerate(image_files):
-        if count < MAX_PARTS:
+        if count < MAX_PARTS:  # only visualize up to MAX_PARTS parts
             outputs[rgb_image].append([Image.open(im) for im in image_files[part]])
 
     return [
@@ -91,19 +116,47 @@ def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_sample
     ]
 
 
-def get_trigger(idx: int, fps: int = 40, oscillate: bool = True):
-    def iter_images(rgb_image: str):
+def get_trigger(
+    idx: int, fps: int = 25, oscillate: bool = True
+) -> Callable[[str], Generator[Image.Image, None, None]]:
+    """
+    Return event listener trigger function for image component to animate image sequence.
+
+    :param idx: index of part to animate from output
+    :param fps: approximate rate at which images should be cycled through in frames per second. Note that the fps cannot
+    be higher than the rate at which images can be returned and rendered. Defaults to 40
+    :param oscillate: if True, animates part in reverse after running from start to end. Defaults to True
+    """
+
+    def iter_images(rgb_image: str) -> Generator[Image.Image, None, None]:
+        """Iterator to yield sequence of images for rendering, based on temp RGB image path"""
+        start_time = time.time()
+
+        def wait_until_next_frame(frame_count: int) -> None:
+            """wait until appropriate time per the specified fps, relative to start time of iteration"""
+            time_to_sleep = max(frame_count / fps - (time.time() - start_time), 0)
+            if time_to_sleep <= 0:
+                print("[WARNING] frames cannot be rendered at the specified FPS due to processing/rendering time.")
+            time.sleep(time_to_sleep)
+
         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]):
+            frame_count = 0
+
+            # iterate forward
             for im in outputs[rgb_image][idx]:
-                time.sleep(1.0 / fps)
+                wait_until_next_frame(frame_count)
                 yield im
+                frame_count += 1
+
+            # iterate in reverse
             if oscillate:
                 for im in reversed(outputs[rgb_image][idx]):
-                    time.sleep(1.0 / fps)
+                    wait_until_next_frame(frame_count)
                     yield im
+                    frame_count += 1
 
         else:
             gr.Error("Could not find any images to load into this module.")
@@ -112,6 +165,9 @@ def get_trigger(idx: int, fps: int = 40, oscillate: bool = True):
 
 
 def clear_outputs():
+    """
+    Remove images from image components.
+    """
     return [gr.update(value=None, visible=(idx == 0)) for idx in range(MAX_PARTS)]
 
 
@@ -119,12 +175,19 @@ with gr.Blocks() as demo:
     gr.Markdown(
         """
     # OPDMulti Demo
-    Upload an image to see its range of motion.
+    We tackle the openable-part-detection (OPD) problem where we identify in a single-view image parts that are openable and their motion parameters. Our OPDFORMER architecture outputs segmentations for openable parts on potentially multiple objects, along with each part’s motion parameters: motion type (translation or rotation, indicated by blue or purple mask), motion axis and origin (see green arrows and points). For each openable part, we predict the motion parameters (axis and origin) in object coordinates along with an object pose prediction to convert to camera coordinates.
+
+    More information about the project, including code, can be found [here](https://3dlg-hcvc.github.io/OPDMulti/).
+
+    Upload an image to see a visualization of its range of motion below. Only the RGB image is needed for the model itself, but the depth image is required as of now for the visualization of motion.
+
+    If you know the intrinsic matrix of your camera, you can specify that here or otherwise use the default matrix which will work with any of the provided examples.
+
+    You can also change the number of samples to define the number of states in the visualization generated.
     """
     )
 
-    # TODO: add gr.Examples
-
+    # inputs
     with gr.Row():
         rgb_image = gr.Image(
             image_mode="RGB", source="upload", type="filepath", label="RGB Image", show_label=True, interactive=True
@@ -133,7 +196,7 @@ with gr.Blocks() as demo:
             image_mode="I;16", source="upload", type="filepath", label="Depth Image", show_label=True, interactive=True
         )
 
-    intrinsics = gr.Dataframe(
+    intrinsic = gr.Dataframe(
         value=[
             [
                 214.85935872395834,
@@ -155,7 +218,7 @@ with gr.Blocks() as demo:
         col_count=(3, "fixed"),
         datatype="number",
         type="numpy",
-        label="Intrinsics matrix",
+        label="Intrinsic matrix",
         show_label=True,
         interactive=True,
     )
@@ -169,6 +232,7 @@ with gr.Blocks() as demo:
         maximum=20,
     )
 
+    # specify examples which can be used to start
     examples = gr.Examples(
         examples=[
             ["examples/59-4860.png", "examples/59-4860_d.png"],
@@ -182,10 +246,12 @@ 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.
+    # output
+    explanation = gr.Markdown(
+        value=f"# Output\nClick on an image to see an animation of the part motion. As of now, only up to {MAX_PARTS} parts can be visualized due to limitations of the visualizer."
+    )
+
     images = [
         gr.Image(type="pil", label=f"Part {idx + 1}", show_download_button=False, visible=(idx == 0))
         for idx in range(MAX_PARTS)
@@ -194,11 +260,11 @@ with gr.Blocks() as demo:
         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)
+    rgb_image.change(clear_outputs, inputs=rgb_image, outputs=images, api_name=False)
+    depth_image.change(clear_outputs, inputs=rgb_image, outputs=images, api_name=False)
 
     submit_btn.click(
-        fn=predict, inputs=[rgb_image, depth_image, intrinsics, num_samples], outputs=images, api_name=False
+        fn=predict, inputs=[rgb_image, depth_image, intrinsic, num_samples], outputs=images, api_name=False
     )
 
 demo.queue(api_open=False)

requirements.txt

@@ -16,3 +16,4 @@ timm==0.9.7
 black==23.9.1
 gradio==3.44.3
 huggingface-hub==0.17.2
+natsort==8.4.0