DOCSEARCHbyAlgolia
Submit Docs

FoundationPose

repository·main·Indexed 25 days ago

https://github.com/nvlabs/foundationpose

A unified foundation model for 6D object pose estimation and tracking. It supports both model-based setups using CAD models and model-free setups using reference images, enabling application to novel objects without fine-tuning. The repository includes instructions for Docker and Conda installation, as well as scripts for running demos on LINEMOD and YCB-Video datasets.

Tokens
1.6K
Snippets
4
Records
6
Agent score
36%

What's inside FoundationPose

  1. Prepare FoundationPose data and weights

    main

    Before running the model, you must download and organize the necessary weights and demo data:

    1. Network Weights: Download weights from the provided Google Drive link and place them in a weights/ folder.
      • For the refiner, use the 2023-10-28-18-33-37 folder.
      • For the scorer, use the 2024-01-11-20-02-45 folder.
    2. Demo Data: Download demo data and extract it into a demo_data/ folder.
    3. Optional Training Data: Large-scale training data is available via Google Drive.
    4. Optional Model-Free Data: For the model-free few-shot version, download the preprocessed reference views.

  2. Run on LINEMOD and YCB-Video Datasets

    main

    Model-Based Version

    Run on LINEMOD or YCB-Video by providing the dataset directory via --linemod_dir or --ycbv_dir. Set --use_reconstructed_mesh 0 for the model-based version.

    python run_linemod.py --linemod_dir /path/to/LINEMOD --use_reconstructed_mesh 0
python run_ycb_video.py --ycbv_dir /path/to/YCB_Video --use_reconstructed_mesh 0

    Model-Free Few-Shot Version

    1. Train Neural Object Field:
    python bundlesdf/run_nerf.py --ref_view_dir /path/to/ref_views --dataset ycbv
    1. Run Pose Estimation: Use --use_reconstructed_mesh 1 and provide the --ref_view_dir.
    python run_ycb_video.py --ycbv_dir /path/to/YCB_Video --use_reconstructed_mesh 1 --ref_view_dir /path/to/ref_views

  3. Install FoundationPose using Conda (Local)

    main

    To install locally using Conda, follow these steps:

    1. Create Environment:
    conda env create -f environment.yml
conda activate foundationpose
    1. Install PyTorch: Match the CUDA build to your machine (e.g., cu124):
    python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
    1. Install PyTorch3D and NVDiffRast: Compile from source. Ensure CUDA_HOME is set to your CUDA installation path.
    export CUDA_HOME=/usr/local/cuda
export PATH="$CUDA_HOME/bin:$PATH"
python -m pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git"
python -m pip install --no-build-isolation "git+https://github.com/NVlabs/nvdiffrast.git"
    1. Install Dependencies and Build Extensions:
    python -m pip install -r requirements.txt
bash build_all_conda.sh
    1. Optional (Model-Free only): Install Kaolin (ensure version matches PyTorch/CUDA).
    conda env create -f environment.yml
conda activate foundationpose
python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
export CUDA_HOME=/usr/local/cuda
export PATH="$CUDA_HOME/bin:$PATH"
python -m pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git"
python -m pip install --no-build-isolation "git+https://github.com/NVlabs/nvdiffrast.git"
python -m pip install -r requirements.txt
bash build_all_conda.sh

  4. Install FoundationPose using Docker (Recommended)

    main

    The recommended way to set up the environment is using Docker.

    1. Pull and tag the image:
    cd docker/
docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose
# Or build from scratch:
docker build --network host -t foundationpose .
    1. Run the container:
    bash docker/run_container.sh
    1. First-time setup: If this is your first launch, you must build the extensions inside the container:
    bash build_all.sh

    Note for RTX 4090 users: Use the custom CUDA 12.1 image instead:

    docker pull shingarey/foundationpose_custom_cuda121:latest

    Then modify the run_container.sh script to use this image.

    cd docker/
docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose
bash docker/run_container.sh
# Inside container:
bash build_all.sh

  5. Run the Model-Based Demo

    main

    Run the default demo using the pre-configured argparse paths. This will perform pose estimation on the first frame and automatically switch to tracking mode for the rest of the video. Visualizations are saved to the debug_dir specified in the arguments.

    python run_demo.py

    To test on other objects (e.g., a driller), modify the paths via argparse arguments.

    python run_demo.py

  6. Parse Camera Parameters from Training Data

    main

    When using the provided training data, use the following logic to convert the camera parameters (including extrinsics and intrinsics) from the JSON format to a standard camera matrix K and world-to-camera transformations.

    import numpy as np
import json

# glcam_in_cvcam maps GL camera to CV camera
glcam_in_cvcam = np.array([[1,0,0,0],
                          [0,-1,0,0],
                          [0,0,-1,0],
                          [0,0,0,1]]).astype(float)

W, H = camera_params["renderProductResolution"]
with open(f'{base_dir}/camera_params/camera_params_000000.json','r') as ff:
    camera_params = json.load(ff)

world_in_glcam = np.array(camera_params['cameraViewTransform']).reshape(4,4).T
cam_in_world = np.linalg.inv(world_in_glcam) @ glcam_in_cvcam
world_in_cam = np.linalg.inv(cam_in_world)

focal_length = camera_params["cameraFocalLength"]
horiz_aperture = camera_params["cameraAperture"][0]
vert_aperture = H / W * horiz_aperture

focal_y = H * focal_length / vert_aperture
focal_x = W * focal_length / horiz_aperture
center_y = H * 0.5
center_x = W * 0.5

fx, fy, cx, cy = focal_x, focal_y, center_x, center_y
K = np.eye(3)
K[0,0] = fx
K[1,1] = fy
K[0,2] = cx
K[1,2] = cy
    import numpy as np
import json

glcam_in_cvcam = np.array([[1,0,0,0],
                          [0,-1,0,0],
                          [0,0,-1,0],
                          [0,0,0,1]]).astype(float)

W, H = camera_params["renderProductResolution"]
with open(f'{base_dir}/camera_params/camera_params_000000.json','r') as ff:
    camera_params = json.load(ff)

world_in_glcam = np.array(camera_params['cameraViewTransform']).reshape(4,4).T
cam_in_world = np.linalg.inv(world_in_glcam) @ glcam_in_cvcam
world_in_cam = np.linalg.inv(cam_in_world)

focal_length = camera_params["cameraFocalLength"]
horiz_aperture = camera_params["cameraAperture"][0]
vert_aperture = H / W * horiz_aperture

focal_y = H * focal_length / vert_aperture
focal_x = W * focal_length / horiz_aperture
center_y = H * 0.5
center_x = W * 0.5

fx, fy, cx, cy = focal_x, focal_y, center_x, center_y
K = np.eye(3)
K[0,0] = fx
K[1,1] = fy
K[0,2] = cx
K[1,2] = cy