Persistent Object Gaussian Splat for Tracking Human and Robot Manipulation of Irregularly Shaped Objects
This repository contains the official implementation for POGS.
Tested on Python 3.10, cuda 11.8, using conda.
- Create conda environment and install relevant packages
conda create --name pogs_env -y python=3.10
conda activate pogs_env
conda install -c "nvidia/label/cuda-11.8.0" cuda-toolkit
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
pip install gsplat --index-url https://docs.gsplat.studio/whl/pt20cu118
pip install warp-lang
-
cumlis required (for global clustering). The best way to install it is with pip:pip install --extra-index-url=https://pypi.nvidia.com cudf-cu11==24.2.* cuml-cu11==24.2.*. -
Install POGS!
git clone https://github.com/uynitsuj/pogs.git --recurse-submodules
cd pogs
python -m pip install -e .
pip install transformers==4.44.0
pip install fast_simplification==0.1.9
pip install numpy==1.26.4
There is also a physical robot component with the UR5 and ZED 2 cameras. To install relevant libraries:
pip install ur_rtde==1.4.2 cowsay opt-einsum pyvista autolab-core
cd ~/pogs/pogs/dependencies/ur5py
pip install -e .
cd ~/pogs/pogs/dependencies/raftstereo
bash download_models.sh
pip install -e .
Contact Graspnet relies on some older library setups, so we couldn't merge everything into 1 conda environment. However, we can make it work by making this separate conda environment and then calling it in a subprocess.
conda deactivate
conda create --name contact_graspnet_env python=3.8
conda activate contact_graspnet_env
conda install -c conda-forge cudatoolkit=11.2
conda install -c conda-forge cudnn=8.2
# If you don't have cuda installed at /usr/local/cuda then you can install on your conda env and run these two lines
conda install -c conda-forge cudatoolkit-dev
export CUDA_HOME=/path/to/anaconda/envs/contact_graspnet_env/bin/nvcc
pip install tensorflow==2.5 tensorflow-gpu==2.5
pip install opencv-python-headless pyyaml pyrender tqdm mayavi
pip install open3d==0.10.0 typing-extensions==3.7.4 trimesh==3.8.12 configobj==5.0.6 matplotlib==3.3.2 pyside2==5.11.0 scikit-image==0.19.0 numpy==1.19.2 scipy==1.9.1 vtk==9.3.1
# if you have cuda installed at /usr/local/cuda run these lines
cd ~/pogs/pogs/dependencies/contact_graspnet
sh compile_pointnet_tfops.sh
# if you have cuda installed on your conda env run these lines
cd ~/pogs/pogs/configs
cp conda_compile_pointnet_tfops.sh ~/pogs/pogs/dependencies/contact_graspnet/
cd ~/pogs/pogs/dependencies/contact_graspnet
sh conda_compile_pointnet_tfops.sh
pip install autolab-core
Download trained models from here and copy them into the checkpoints/ folder.
Download the test data from here and copy them them into the test_data/ folder.
Before training/tracking POGS, make sure wrist mounted camera and third-person view camera are calibrated. We use an Aruco marker for the calibration
conda activate pogs_env
cd ~/pogs/pogs/scripts
python calibrate_cameras.py
Script used to perform hemisphere capture with robot on tabletop scene. We used manual trajectory but you can also put the robot in "teach" mode to capture trajectory.
conda activate pogs_env
cd ~/pogs/pogs/scripts
python scene_capture.py --scene DATA_NAME
Script used to train the POGS for 4000 steps
conda activate pogs_env
ns-train pogs --data /path/to/data/folder
Once the POGS has completed training, there are N steps to then actually define/save the object clusters.
- Hit the cluster scene button.
- It will take 10-20 seconds, but then after, you should see your objects as specific clusters. If not, hit Toggle RGB/Cluster and try to cluster the scene again but change the Cluster Eps (lower normally works better).
- Once you have your scene clustered, hit Toggle RGB/Cluster.
- Then, hit Click and click on your desired object (green ball will appear on object).
- Hit Crop to Click, and it should isolate the object.
- A draggable coordinate frame will pop up to indicate the object's origin, drag it to where you want it to be. (For experiments, this was what we used to align for object reset or tool servoing)
- Hit Add Crop to Group List
- Repeat steps 4-7 for all objects in scene
- Hit View Crop Group List Once you have trained the POGS, make sure you have the config file and checkpoint directory from the terminal saved.
Script for letting you use a POGS to track an object online and grasp it.
conda activate pogs_env
python ~/pogs/pogs/scripts/track_main_online_demo.py --config_path /path/to/config/yml
