Add Docker Orbbec workflow and live camera grasp demo

docker/build.sh

@@ -1,2 +1,14 @@
+#!/bin/bash
+set -euo pipefail
+
 VER=1.0
-docker build -f docker/graspgen_cuda121.dockerfile --progress=plain . --network=host -t graspgen:$VER -t graspgen:latest
+
+BUILD_ARGS=()
+if [[ -n "${PYORBBECSDK2_WHL:-}" ]]; then
+	BUILD_ARGS+=(--build-arg "PYORBBECSDK2_WHL=${PYORBBECSDK2_WHL}")
+	echo "Using PYORBBECSDK2_WHL=${PYORBBECSDK2_WHL}"
+else
+	echo "PYORBBECSDK2_WHL not set; building without pyorbbecsdk2 preinstall"
+fi
+
+docker build -f docker/graspgen_cuda121.dockerfile --progress=plain . --network=host "${BUILD_ARGS[@]}" -t graspgen:$VER -t graspgen:latest

docker/graspgen_cuda121.dockerfile

@@ -1,6 +1,8 @@
 # Build base image
 FROM nvcr.io/nvidia/pytorch:23.07-py3 AS base
 
+ARG PYORBBECSDK2_WHL=""
+
 # tmux is for debugging, osmesa is for rendering. Put all apt-get installs in this line
 RUN apt update && apt-get install -y tmux libosmesa6-dev
 
@@ -24,7 +26,7 @@ RUN pip install pyrender==0.1.45 pyglet==2.1.6 && pip install PyOpenGL==3.1.5
 
 # Install pointnet2 modules
 COPY pointnet2_ops pointnet2_ops
-RUN pip install ./pointnet2_ops
+RUN pip install --no-build-isolation ./pointnet2_ops
 
 # Diffusion dependencies
 RUN pip install diffusers==0.11.1 timm==1.0.15
@@ -57,4 +59,11 @@ RUN cd /install/Manifold/build && cmake .. -DCMAKE_BUILD_TYPE=Release
 RUN cd /install/Manifold/build && make
 ENV PATH="${PATH}:/install/Manifold/build/"
 
+RUN if [ -n "$PYORBBECSDK2_WHL" ]; then \
+			echo "Installing pyorbbecsdk2 from $PYORBBECSDK2_WHL" && \
+			pip install "$PYORBBECSDK2_WHL"; \
+		else \
+			echo "Skipping pyorbbecsdk2 install (PYORBBECSDK2_WHL not set)"; \
+		fi
+
 WORKDIR /code/

docker/run.sh

@@ -149,6 +149,9 @@ docker run \
   -e NVIDIA_DISABLE_REQUIRE=1 \
   -e NVIDIA_DRIVER_CAPABILITIES=all \
   --device /dev/dri \
+    --device /dev/bus/usb:/dev/bus/usb \
+    --device-cgroup-rule='c 189:* rmw' \
+    -v /run/udev:/run/udev:ro \
   -it \
   -e DISPLAY \
   $VOLUME_MOUNTS \
@@ -157,5 +160,5 @@ docker run \
   --shm-size 40G \
   graspgen:latest \
   /bin/bash \
-  -c "cd /code/ && pip install -e . && bash" \
+    -c "cd /code/ && pip install -e . && bash"
 xhost -local:root

grasp_gen/utils/point_cloud_utils.py

@@ -36,6 +36,10 @@ def knn_points(X: torch.Tensor, K: int, norm: int):
         idxs: (N, K) tensor containing indices of the K nearest neighbors.
     """
     N, _ = X.shape
+    if N == 0:
+        raise ValueError("Point cloud is empty")
+
+    effective_k = min(K, max(1, N - 1))
 
     # Compute pairwise squared Euclidean distances
     dist_matrix = torch.cdist(X, X, p=norm)  # (N, N)
@@ -45,7 +49,7 @@ def knn_points(X: torch.Tensor, K: int, norm: int):
     dist_matrix.masked_fill_(self_mask, float("inf"))  # Set self-distances to inf
 
     # Get the indices of the K-nearest neighbors
-    dists, idxs = torch.topk(dist_matrix, K, dim=1, largest=False)
+    dists, idxs = torch.topk(dist_matrix, effective_k, dim=1, largest=False)
 
     return dists, idxs
 
@@ -74,11 +78,26 @@ def point_cloud_outlier_removal(
     obj_pc = obj_pc.float()
     obj_pc = obj_pc.unsqueeze(0)
 
+    num_points = obj_pc.shape[1]
+    if num_points <= max(1, K):
+        removed_pc = obj_pc[0].new_empty((0, 3))
+        logger.info(
+            f"Skipping outlier removal for small point cloud ({num_points} points)"
+        )
+        return obj_pc[0].view(-1, 3), removed_pc
+
     nn_dists, _ = knn_points(obj_pc[0], K=K, norm=1)
 
     mask = nn_dists.mean(1) < threshold
     filtered_pc = obj_pc[0, mask]
     removed_pc = obj_pc[0][~mask]
+
+    if filtered_pc.shape[0] == 0:
+        logger.info(
+            "Outlier removal removed all points; falling back to original point cloud"
+        )
+        return obj_pc[0].view(-1, 3), obj_pc[0].new_empty((0, 3))
+
     filtered_pc = filtered_pc.view(-1, 3)
     removed_pc = removed_pc.view(-1, 3)
 
@@ -119,11 +138,37 @@ def point_cloud_outlier_removal_with_color(
     obj_pc_color = obj_pc_color.float()
     obj_pc_color = obj_pc_color.unsqueeze(0)
 
+    num_points = obj_pc.shape[1]
+    if num_points <= max(1, K):
+        removed_pc = obj_pc[0].new_empty((0, 3))
+        removed_pc_color = obj_pc_color[0].new_empty((0, 3))
+        logger.info(
+            f"Skipping outlier removal for small point cloud ({num_points} points)"
+        )
+        return (
+            obj_pc[0].view(-1, 3),
+            removed_pc,
+            obj_pc_color[0].view(-1, 3),
+            removed_pc_color,
+        )
+
     nn_dists, _ = knn_points(obj_pc[0], K=K, norm=1)
 
     mask = nn_dists.mean(1) < threshold
     filtered_pc = obj_pc[0, mask]
     removed_pc = obj_pc[0][~mask]
+
+    if filtered_pc.shape[0] == 0:
+        logger.info(
+            "Outlier removal removed all points; falling back to original point cloud/colors"
+        )
+        return (
+            obj_pc[0].view(-1, 3),
+            obj_pc[0].new_empty((0, 3)),
+            obj_pc_color[0].view(-1, 3),
+            obj_pc_color[0].new_empty((0, 3)),
+        )
+
     filtered_pc = filtered_pc.view(-1, 3)
     removed_pc = removed_pc.view(-1, 3)