Docker support and various fixes to get it working

.dockerignore

@@ -0,0 +1,2 @@
+dataset
+docker

README.md

@@ -22,6 +22,16 @@ We strongly suggest to install it in either a
 [venv](https://docs.python.org/fr/3/library/venv.html) or a
 [conda environment](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html).
 
+### Example with Docker
+
+```
+git clone --branch dev --recurse-submodules https://github.com/agimus-project/happypose.git
+cd happypose
+docker build -f docker/Dockerfile -t happypose:latest .
+source ./docker/to_source.sh
+happypose_docker
+```
+
 ### Example with conda
 
 ```
@@ -107,11 +117,17 @@ python -m happypose.toolbox.utils.download --cosypose_models \
 ```
 
 ```
-python -m happypose.toolbox.utils.download --bop_dataset ycbv
+hf download bop-benchmark/ycbv \
+            --local-dir ./dataset/bop_datasets/ycbv \
+            --repo-type=dataset \
+            ycbv_base.zip ycbv_models.zip ycbv_test_all.zip ycbv_train_pbr.zip
+cd ./dataset/bop_datasets/ycbv
+7z e ycbv_base.zip
+7z x ycbv_base.zip ycbv_models.zip ycbv_test_all.zip
 ```
 
 ```
-python -m happypose.toolbox.utils.download --test-results
+python -m happypose.toolbox.utils.download --ycbv_tests
 ```
 
-The tests take much longer in this case.
+Run the tests again using pytest as above, this time they will take much longer to complete.

docker/Dockerfile

@@ -0,0 +1,46 @@
+# syntax=docker/dockerfile:1.7-labs
+FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
+
+# Install system-wide dependencies
+RUN apt-get update && \
+    apt install -yq software-properties-common && \
+    add-apt-repository -y ppa:mozillateam/ppa && \
+    apt update && \
+    apt install -yq \
+        curl \
+        nano \
+        git \
+        python3 \
+        python3-pip \
+        p7zip-full \
+        ffmpeg libsm6 libxext6 \
+        nvidia-container-runtime \
+        libgl1 libegl1-mesa libgles2-mesa \
+        firefox=153.0.1+build2-0ubuntu0.22.04.1~mt1 \
+        && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Change default python version
+RUN ln -s /usr/bin/python3 /usr/bin/python
+
+# Install uv
+RUN pip install uv==0.8.24
+
+# Tmp copy of repo's root
+WORKDIR /happypose
+COPY . .
+
+# Install pypi deps
+RUN uv sync --extra pypi --extra cu124
+
+# Install geckodriver
+RUN curl -LO "https://github.com/mozilla/geckodriver/releases/download/v0.37.1/geckodriver-v0.37.1-linux64.tar.gz"
+RUN mkdir -p /opt/geckodriver && tar -xvf geckodriver-v0.37.1-linux64.tar.gz -C /opt/geckodriver
+ENV PATH="$PATH:/opt/geckodriver"
+
+# Create unpriviledged user
+RUN useradd -ms /bin/bash happypose
+USER happypose
+
+CMD ["/bin/bash"]

docker/to_source.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+happypose_docker() {
+  echo "starting happypose docker"
+  xhost +local:docker;
+  mkdir -p $HOME/happypose/dataset
+  docker run -it --rm -d \
+    --name="happypose_dev" \
+    --runtime=nvidia \
+    --workdir $HOME/happypose \
+    -v /etc/localtime:/etc/localtime:ro \
+    -v /dev/input:/dev/input \
+    --shm-size=12G \
+    --net=host \
+    --add-host happypose_dev:127.0.0.1 \
+    --hostname=happypose_dev \
+    --privileged=true \
+    --env=QT_X11_NO_MITSHM=1 \
+    --device=/dev/dri:/dev/dri \
+    -v /tmp/.X11-unix:/tmp/.X11-unix:rw \
+    -v /etc/localtime:/etc/localtime:ro \
+    -v "./:$HOME/happypose" \
+    -e HAPPYPOSE_DATA_DIR="$HOME/happypose/dataset" \
+    happypose:latest
+
+    happypose_docker_attach;
+}
+
+happypose_docker_attach() {
+  docker exec -it -e "COLUMNS=$COLUMNS" -e "LINES=$LINES" happypose_dev /bin/bash -c "source /happypose/.venv/bin/activate && bash"
+}

happypose/pose_estimators/megapose/evaluation/bop.py

@@ -104,7 +104,7 @@ def convert_results_to_coco(results_path, out_json_path, detection_method):
 def convert_results_to_bop(
     results_path: Path, out_csv_path: Path, method: str, use_pose_score: bool = True
 ):
-    predictions = torch.load(results_path)["predictions"]
+    predictions = torch.load(results_path, weights_only=False)["predictions"]
     predictions = predictions[method]
     if method == "coarse":
         predictions = get_best_coarse_predictions(predictions)

happypose/pose_estimators/megapose/evaluation/evaluation.py

@@ -49,6 +49,8 @@
 
 logger = get_logger(__name__)
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 def generate_save_key(detection_type: str, coarse_estimation_type: str) -> str:
     return f"{detection_type}+{coarse_estimation_type}"
@@ -125,7 +127,7 @@ def run_eval(
     if cfg.inference.detection_type == "detector":
         assert cfg.detector_run_id is not None
         detector_model = happypose.toolbox.inference.utils.load_detector(
-            cfg.detector_run_id,
+            cfg.detector_run_id, device
         )
     elif cfg.inference.detection_type in ["gt", "exte"]:
         detector_model = None

happypose/pose_estimators/megapose/inference/pose_estimator.py

@@ -333,6 +333,7 @@ def forward_coarse_model(
         - Scores them using the coarse model.
         """
         start_time = time.time()
+        device = observation.images.device
 
         happypose.toolbox.inference.types.assert_detections_valid(detections)
 
@@ -342,6 +343,9 @@ def forward_coarse_model(
         B = len(detections)
         M = self._SO3_grid.shape[0]
 
+        if SO3_grid.device != device:
+            SO3_grid = SO3_grid.to(device)
+
         # Add M rows for each row in detections
         df = detections.infos
         df_concat = []
@@ -359,7 +363,6 @@ def forward_coarse_model(
         ids = torch.arange(len(df_hypotheses))
         ds = TensorDataset(ids)
         dl = DataLoader(ds, batch_size=bsz_images)
-        device = observation.images.device
 
         images_crop_list = []
         renders_list = []

happypose/pose_estimators/megapose/scripts/run_inference_on_example.py

@@ -45,7 +45,7 @@ def run_inference(
     model_info: Dict,
     observation: ObservationTensor,
     detections: DetectionsType,
-) -> None:
+):
     observation.to(device)
 
     logger.info("Running inference.")
@@ -104,6 +104,9 @@ def run_inference(
         output = run_inference(pose_estimator, model_info, observation, detections)
         save_predictions(output, example_dir)
 
+        # TODO: tmp solution to megapose stalling
+        del pose_estimator
+
     if args.vis_detections:
         make_detections_visualization(rgb, detections, example_dir)
 

pyproject.toml

@@ -128,6 +128,9 @@ dependency-metadata = [
 [tool.uv.build-backend]
 module-root = ""
 
+[tool.uv.extra-build-dependencies]
+visdom = ["setuptools<82.0.0"]
+
 [[tool.uv.index]]
 explicit = true
 name = "pytorch-cpu"