DOCSEARCHbyAlgolia
Submit Docs

Sapiens2: High-Resolution Transformers for Human-Centric Vision

repository·main·Indexed 21 days ago

https://github.com/facebookresearch/sapiens2

Sapiens2 is a family of high-resolution transformers pretrained on 1 billion human images, designed for human-centric vision tasks. It supports pose estimation (308 keypoints), body-part segmentation (29 classes), surface normal estimation, 3D pointmap generation, and human matting. The library provides pretrained checkpoints ranging from 0.1B to 5B parameters and supports standalone execution using only torch and safetensors.

Tokens
24.6K
Snippets
92
Records
108
Agent score
74%

What's inside Sapiens2

  1. Available Vision Tasks in Sapiens2

    main

    Sapiens2 supports several human-centric vision tasks. Detailed inference and training guides for each task are available in the docs/ directory:

    • Pose Estimation: 308 whole-body keypoints. See docs/POSE.md.
    • Body-Part Segmentation: 29 body parts. See docs/SEG.md.
    • Surface Normal Estimation: Per-pixel normals. See docs/NORMAL.md.
    • Pointmap Estimation: Per-pixel 3D points. See docs/POINTMAP.md.
    • Human Matting: Alpha matte + foreground. See docs/MATTING.md.

  2. Prepare data for Surface Normal Estimation training

    main

    To train Sapiens2 for surface normal estimation, organize your dataset according to the following structure and requirements:

    1. Environment Variable

    Set the $DATA_ROOT environment variable to the root directory of your dataset:

    export DATA_ROOT=/path/to/your/normal_data

    2. Directory Structure

    Your dataset must follow this layout:

    $DATA_ROOT/
├── images/                  # input RGB images
├── normals/                 # ground-truth normals (.npy or .png)
└── annotations/
    └── train.json           # image/normal pair index

    3. Normal Encoding

    Surface normals must be 3-channel (x, y, z) unit vectors per pixel. They can be provided as:

    • .npy arrays
    • RGB images using the standard remapping: n = 2*rgb - 1

    For reference implementations of dataset classes, see normal_metasim_dataset.py, normal_render_people_body_dataset.py, or normal_thuman_dataset.py in the repository.

    export DATA_ROOT=/path/to/your/normal_data

  3. Install Sapiens2

    main

    To install the full Sapiens2 repository, follow these steps:

    1. Clone the repository:

      git clone https://github.com/facebookresearch/sapiens2.git
cd sapiens2
export SAPIENS_ROOT=$(pwd)

    2. Install via pip: Requires Python $\ge$ 3.12 and PyTorch $\ge$ 2.7.

      pip install -e .

    3. Checkpoints: Download checkpoints (refer to docs/MODEL_ZOO.md) and place them under $SAPIENS_CHECKPOINT_ROOT. By default, this is ~/sapiens2_host.

    git clone https://github.com/facebookresearch/sapiens2.git
cd sapiens2
export SAPIENS_ROOT=$(pwd)
pip install -e .

  4. Download Sapiens2 Surface Normal checkpoints

    main

    To use Sapiens2 for surface normal estimation, download the model checkpoints from HuggingFace and place them in the directory specified by the $SAPIENS_CHECKPOINT_ROOT environment variable. By default, this directory is ~/sapiens2_host.

    Available models and their expected paths:

    • Sapiens2-0.4B: $SAPIENS_CHECKPOINT_ROOT/normal/sapiens2_0.4b_normal.safetensors
    • Sapiens2-0.8B: $SAPIENS_CHECKPOINT_ROOT/normal/sapiens2_0.8b_normal.safetensors
    • Sapiens2-1B: $SAPIENS_CHECKPOINT_ROOT/normal/sapiens2_1b_normal.safetensors
    • Sapiens2-5B: $SAPIENS_CHECKPOINT_ROOT/normal/sapiens2_5b_normal.safetensors

  5. Quick Start: Run a pretrained backbone forward pass

    main

    You can run a pretrained Sapiens2 backbone using only torch and safetensors. This approach is useful for getting dense backbone features from an image.

    Requirements:

    • torch
    • safetensors
    • An RGB image (ImageNet normalization is recommended).

    Note: The img_size parameter should be provided as (H, W).

    import os
import torch
from safetensors.torch import load_file
from sapiens.backbones.standalone.sapiens2 import Sapiens2

# Build the model and load a pretrained checkpoint
model = Sapiens2(arch="sapiens2_1b", img_size=(1024, 768), patch_size=16).eval().cuda()  # img_size is (H, W)
ckpt = os.path.expanduser("~/sapiens2_host/pretrain/sapiens2_1b_pretrain.safetensors")
model.load_state_dict(load_file(ckpt))

# Forward pass on a single image (RGB; ImageNet normalization recommended)
x = torch.randn(1, 3, 1024, 768).cuda()
with torch.no_grad():
    features = model(x)[0]  # dense backbone features

  6. Download Sapiens2 Pose Models

    main

    Sapiens2 provides 308-keypoint top-down pose estimation (including face, hand, and foot keypoints) following the Sociopticon keypoint format.

    To use the models, download the checkpoints from HuggingFace and place them in the directory specified by the $SAPIENS_CHECKPOINT_ROOT environment variable (defaults to ~/sapiens2_host).

    ModelCheckpoint Path
    Sapiens2-0.4B$SAPIENS_CHECKPOINT_ROOT/pose/sapiens2_0.4b_pose.safetensors
    Sapiens2-0.8B$SAPIENS_CHECKPOINT_ROOT/pose/sapiens2_0.8b_pose.safetensors
    Sapiens2-1B$SAPIENS_CHECKPOINT_ROOT/pose/sapiens2_1b_pose.safetensors
    Sapiens2-5B$SAPIENS_CHECKPOINT_ROOT/pose/sapiens2_5b_pose.safetensors

  7. Launch Surface Normal training via node.sh

    main

    Training is initiated using the node.sh script located in the dense directory.

    Execution Command

    cd $SAPIENS_ROOT/sapiens/dense
./scripts/normal/train/sapiens2_1b/node.sh

    Script Configuration

    You can modify node.sh to adjust the following parameters:

    • DEVICES: GPU IDs to use (e.g., 0,1,2,3,4,5,6,7).
    • TRAIN_BATCH_SIZE_PER_GPU: The batch size for each individual GPU.
    • mode: Set to 'multi-gpu' for production or 'debug' for a single-GPU dry-run.
    • LOAD_FROM: Path to a checkpoint to initialize weights.
    • RESUME_FROM: Path to a checkpoint to resume training.

    Output Location

    Checkpoints and logs are saved to: Outputs/normal/train/sapiens2_1b_normal_metasim_render_people-1024x768/node/<timestamp>/

    cd $SAPIENS_ROOT/sapiens/dense
./scripts/normal/train/sapiens2_1b/node.sh

  8. Run Human Image Matting inference

    main

    Sapiens2 provides per-pixel alpha matting for human subjects, predicting a soft foreground mask (alpha in [0, 1]) and a pre-multiplied foreground RGB output.

    To run inference on a demo dataset, navigate to the dense directory and execute the matting script:

    cd $SAPIENS_ROOT/sapiens/dense
./scripts/demo/matting.sh

    Configuration

    You can customize the behavior by editing the ./scripts/demo/matting.sh script. The following variables are configurable:

    • INPUT: Path to your image directory (default: ../../demo/data).
    • OUTPUT: Directory where visualizations will be saved.
    • MODEL_NAME: The model size to use (default: sapiens2_1b).
    • JOBS_PER_GPU: Number of parallel jobs per GPU (default: 3).
    • GPU_IDS: List of GPU IDs to use (default: 0-7).

    Outputs

    For every input image, the script generates:

    • A side-by-side visualization containing: [input | alpha matte | foreground on chroma green].
    • If --save_pred is enabled (default in the script), it also writes a <name>_alpha.npy file containing the raw alpha mask as float32 values in the range [0, 1].
    cd $SAPIENS_ROOT/sapiens/dense
./scripts/demo/matting.sh

  9. Download and Organize Sapiens2 Segmentation Checkpoints

    main

    Download the segmentation checkpoints from HuggingFace and place them in the directory specified by the $SAPIENS_CHECKPOINT_ROOT environment variable (defaults to ~/sapiens2_host).

    Required directory structure:

    • $SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_0.4b_seg.safetensors
    • $SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_0.8b_seg.safetensors
    • $SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_1b_seg.safetensors
    • $SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_5b_seg.safetensors
    | Model | Checkpoint Path |
|-------|-----------------|
| Sapiens2-0.4B | `$SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_0.4b_seg.safetensors` |
| Sapiens2-0.8B | `$SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_0.8b_seg.safetensors` |
| Sapiens2-1B   | `$SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_1b_seg.safetensors` |
| Sapiens2-5B   | `$SAPIENS_CHECKPOINT_ROOT/seg/sapiens2_5b_seg.safetensors` |

  10. Download Pointmap model checkpoints

    main

    Download the required checkpoints from HuggingFace and place them in the directory specified by the $SAPIENS_CHECKPOINT_ROOT environment variable (defaults to ~/sapiens2_host).

    Available models:

    • Sapiens2-0.4B: $SAPIENS_CHECKPOINT_ROOT/pointmap/sapiens2_0.4b_pointmap.safetensors
    • Sapiens2-0.8B: $SAPIENS_CHECKPOINT_ROOT/pointmap/sapiens2_0.8b_pointmap.safetensors
    • Sapiens2-1B: $SAPIENS_CHECKPOINT_ROOT/pointmap/sapiens2_1b_pointmap.safetensors
    • Sapiens2-5B: $SAPIENS_CHECKPOINT_ROOT/pointmap/sapiens2_5b_pointmap.safetensors

  11. Run Sapiens2 Body-Part Segmentation Inference

    main

    To run inference on a set of images, use the provided shell script located in the sapiens/dense directory.

    Execution

    cd $SAPIENS_ROOT/sapiens/dense
./scripts/demo/seg.sh

    Configuration

    Before running, edit ./scripts/demo/seg.sh to adjust the following variables:

    • INPUT: Path to your image directory (default: ../../demo/data).
    • OUTPUT: Directory where visualizations will be saved.
    • MODEL_NAME: Uncomment the specific model size you wish to use (0.4B, 0.8B, 1B, or 5B).
    • JOBS_PER_GPU: Number of parallel jobs per GPU (default: 3).
    • GPU_IDS: List of GPU IDs to use (default: 0–7).

    Outputs

    • Visualization images: Color-coded class overlays.
    • .npy files: Contains raw class probabilities and foreground masks. These files are intended for use in downstream tasks like NORMAL, ALBEDO, or POINTMAP generation.
    cd $SAPIENS_ROOT/sapiens/dense
./scripts/demo/seg.sh

  12. Launch Sapiens2 training for Human Image Matting

    main

    To start single-node multi-GPU training for human image matting, navigate to the dense directory and execute the node script.

    Before launching, you must edit the configuration file located at sapiens/dense/configs/matting/gss_p3m_metasim/sapiens2_1b_matting_gss_p3m_metasim-1024x768.py to update the dataset_paths and pretrained_checkpoint to point to your local data and backbone checkpoint.

    Training outputs are saved to: Outputs/matting/train/<MODEL>/node/<timestamp>/.

    cd $SAPIENS_ROOT/sapiens/dense
./scripts/matting/train/sapiens2_1b/node.sh