move model to huggingface

README.md

@@ -19,7 +19,7 @@ Please refer there for more information about the proect and implementation.
 
 ## Installation
 
-## Requirements
+### Requirements
 
 For the docker build, you will just need docker in order to build and run the container, else you will need
 
@@ -30,18 +30,17 @@ For the docker build, you will just need docker in order to build and run the co
 
 A full list of other packages can be found in the Dockerfile, or in `Open3D/util/install_deps_ubuntu.sh`.
 
-**BEFORE BUILDING** as of writing, you will need to copy the model file manually to `.data/models/motion_state_pred_opdformerp_rgb.pth`
-in the repository. This step must occur before the docker build, or if building locally then before running. Future
-work will make this step no longer required.
+The model file can currently be found [here](https://huggingface.co/3dlg-hcvc/opdmulti-motion-state-rgb-model) and is
+downloaded as part of the demo code.
 
-### Docker (preferred)
+### Docker Build (preferred)
 
 To build the docker container, run
 ```
 docker build -f Dockerfile -t opdmulti-demo .
 ```
 
-### Local
+### Local Build
 
 To setup the environment, run the following (recommended in a virtual environment):
 ```

app.py

@@ -3,11 +3,12 @@ import re
 import shutil
 import time
 from types import SimpleNamespace
-from typing import Any, Callable, Generator, Optional
+from typing import Any, Callable, Generator
 
 import gradio as gr
 import numpy as np
 from detectron2 import engine
+from huggingface_hub import hf_hub_download
 from natsort import natsorted
 from PIL import Image
 
@@ -20,7 +21,7 @@ SCORE_THRESHOLD = 0.8
 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",
+    model={"repo_id": "3dlg-hcvc/opdmulti-motion-state-rgb-model", "filename": "pytorch_model.pth"},
     input_format="RGB",
     output=".output",
     cpu=True,
@@ -87,6 +88,8 @@ def predict(rgb_image: str, depth_image: str, intrinsic: np.ndarray, num_samples
         return [None] * 5
 
     # run model
+    weights_path = hf_hub_download(repo_id=ARGS.model["repo_id"], filename=ARGS.model["filename"])
+    ARGS.model = weights_path
     cfg = setup_cfg(ARGS)
     engine.launch(
         main,
@@ -117,7 +120,7 @@ def predict(rgb_image: str, depth_image: str, intrinsic: np.ndarray, num_samples
 
 
 def get_trigger(
-    idx: int, fps: int = 25, oscillate: bool = True
+    idx: int, fps: int = 15, oscillate: bool = True
 ) -> Callable[[str], Generator[Image.Image, None, None]]:
     """
     Return event listener trigger function for image component to animate image sequence.
@@ -260,8 +263,8 @@ 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=rgb_image, outputs=images, api_name=False)
-    depth_image.change(clear_outputs, inputs=rgb_image, outputs=images, api_name=False)
+    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, intrinsic, num_samples], outputs=images, api_name=False

inference.py

@@ -44,7 +44,6 @@ from visualization import (
     generate_rotation_visualization,
     generate_translation_visualization,
     batch_trim,
-    create_gif,
 )
 
 # import based on torch version. Required for model loading. Code is taken from fvcore.common.checkpoint, in order to