基于sapiens进行人体部位分割

七崽爱吃小饼干2025/09/04阅读 7

一、简介

Sapiens 是 Meta Reality Labs 开发的一系列以人为中心的视觉模型。
Sapiens 旨在处理四个基本的以人为中心的视觉任务，包括 2D 姿态估计、身体部位分割、深度估计和表面法线预测。
该模型使用一个庞大的数据集 Humans-300M 进行预训练，其中包含 3 亿张真实世界中的人类图像。这些图像经过过滤以确保质量，并使用自监督预训练方法，如 Masked Autoencoder（MAE），在 1K 分辨率的图像上进行处理。

模型特点

通用性强：具备很强的泛化能力，能在不同场景中表现良好，即使在标注数据稀缺或完全合成的情况下，也能对自然数据表现出出色的泛化能力。
适应性高：只需对在大量自然人类图像上预训练的模型进行少量微调，就可用于特定任务。
高保真度：输出高分辨率、精确的结果，尤其适合人类图像的生成和处理。

二、安装

1.克隆代码

github仓库

克隆代码

shell

git clone https://github.com/facebookresearch/sapiens.git

pwd获取当前项目的根目录

shell

pwd

配置根目录环境变量

shell

export SAPIENS_ROOT=/path/to/sapiens

2.环境配置

项目提供了完整的环境配置脚本，会创建名为sapiens的conda环境，并安装所需要的包

运行命令

bash

cd ./_install
sh conda.sh

3.从`hugging-face` 下载模型权重

链接

3.1选择格式

sapiens提供了三种checkpoints

original：可以根据您的用例和推理微调权重。
torchscript: （仅推理）权重移植到 torchscript。
bfloat16: （仅推理）用于大规模处理，权重移植到 bfloat16（仅 A100 GPU + pytorch-2.3）。第三种格式是专门给数据处理型GPU用的，而且我们想微调模型，所以这里我们用original

3.2选择模型规模

sapiens提供了三种模型规模，不同的规模对显存的要求也不同

sapiens-seg-0.3b
sapiens-seg-0.6b
sapiens-seg-1b

各模型的 GPU 显存要求（最低配置）

模型	推理（1024x1024）	微调（1024x1024，全参数）	微调（1024x1024，LoRA 参数高效微调）
sapiens-seg-0.3b	8-12GB	16-20GB	8-12GB
sapiens-seg-0.6b	12-16GB	24-32GB	12-16GB
sapiens-seg-1b	16-24GB	32-40GB	16-24GB

这里我们选择使用6亿参数模型也就是[sapiens-seg-0.6b](https://huggingface.co/facebook/sapiens-seg-0.6b)

3.3下载

这个暂时没有提供transformer的下载方式，所以需要我们自己下载

3.3.1用cli工具下载

我们可以通过使用 git lfs下载模型

bash

# 安装 Git LFS
sudo apt-get update
sudo apt-get install git-lfs

# 初始化lfs
git lfs install

# 下载模型
git clone https://huggingface.co/facebook/sapiens-seg-0.6b

3.3.2 手动下载

切到files and versions直接下载文件

下载后的模型按照要求存放到对应的目录

4. 安装seg的环境要求

seg目录下有setup.py文件，进入该目录执行命令：

bash

pip install -e .

三、执行推理脚本

推理脚本的demo位置如下

在这个脚本中需要配置一些重要的配置项（checkpoint的选择与路径、输入输出目录以及GPU配置等）：

shell

SAPIENS_CHECKPOINT_ROOT=/home/Minda2024-9/liujingmin/sapiens/sapiens

#----------------------------set your input and output directories----------------------------------------------
INPUT='/home/Minda2024-9/liujingmin/sapiens/sapiens/seg/data/input'
OUTPUT="/home/Minda2024-9/liujingmin/sapiens/sapiens/seg/data/output"

#--------------------------MODEL CARD---------------
# MODEL_NAME='sapiens_0.3b'; CHECKPOINT=$SAPIENS_CHECKPOINT_ROOT/seg/checkpoints/sapiens_0.3b/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194.pth
 MODEL_NAME='sapiens_0.6b'; CHECKPOINT=$SAPIENS_CHECKPOINT_ROOT/seg/checkpoints/sapiens_0.6b/sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178.pth
#MODEL_NAME='sapiens_1b'; CHECKPOINT=$SAPIENS_CHECKPOINT_ROOT/seg/checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151.pth

DATASET='goliath'
MODEL="${MODEL_NAME}_${DATASET}-1024x768"
CONFIG_FILE="configs/sapiens_seg/${DATASET}/${MODEL}.py"
OUTPUT=$OUTPUT/$MODEL_NAME

##-------------------------------------inference-------------------------------------
RUN_FILE='demo/demo_seg_vis.py'

## number of inference jobs per gpu, total number of gpus and gpu ids
JOBS_PER_GPU=1; TOTAL_GPUS=1; VALID_GPU_IDS=(0 1 2 3 4 5 6 7)