访问官网
Github
Huggingface
论文
英文 | 数据增强API文档 | 知乎专栏 | Pypi下载量
最新消息 👇👇
目前我的大部分时间都集中在大语言模型/视觉语言模型推理上。请查看 📖大语言模型推理精选 
、📖Stable Diffusion推理精选 
和 📖CUDA学习笔记 
以获取更多详情。
🤗 简介
torchlm 旨在为人脸关键点检测构建一个高级流程，支持训练、评估、导出、推理(Python/C++)以及100多种数据增强，可以通过pip轻松安装。
👋 核心特性
- 用于训练和推理的高级流程。
- 提供30多种原生关键点数据增强。
- 可以一行代码绑定来自torchvision和albumentations的80多种变换。
- 支持PIPNet、YOLOX、ResNet、MobileNet和ShuffleNet用于人脸关键点检测。
🆕 新特性
- [2022/03/08]: 添加 PIPNet: Towards Efficient Facial Landmark Detection in the Wild, CVPR2021
🔥🔥性能(@NME)
| 模型 | 骨干网络 | 头部 | 300W | COFW | AFLW | WFLW | 下载 |
|---|---|---|---|---|---|---|---|
| PIPNet | MobileNetV2 | 热图+回归+NRM | 3.40 | 3.43 | 1.52 | 4.79 | 链接 |
| PIPNet | ResNet18 | 热图+回归+NRM | 3.36 | 3.31 | 1.48 | 4.47 | 链接 |
| PIPNet | ResNet50 | 热图+回归+NRM | 3.34 | 3.18 | 1.44 | 4.48 | 链接 |
| PIPNet | ResNet101 | 热图+回归+NRM | 3.19 | 3.08 | 1.42 | 4.31 | 链接 |
🛠️安装
你可以直接从pypi安装torchlm。
pip install torchlm>=0.1.6.10 # 或安装最新的pypi版本 `pip install torchlm`
pip install torchlm>=0.1.6.10 -i https://pypi.org/simple/ # 或使用'-i'从特定pypi镜像安装
如果你想要最新的torchlm并以可编辑模式安装，可以从源代码安装，使用-e。
git clone --depth=1 https://github.com/DefTruth/torchlm.git
cd torchlm && pip install -e .
🌟🌟数据增强
torchlm 提供了30多种原生地标数据增强方法，并可以与来自torchvision和albumentations的80多种变换进行绑定。地标的布局格式为xy，形状为(N, 2)。
直接使用torchlm中的30多种原生变换
import torchlm
transform = torchlm.LandmarksCompose([
torchlm.LandmarksRandomScale(prob=0.5),
torchlm.LandmarksRandomMask(prob=0.5),
torchlm.LandmarksRandomBlur(kernel_range=(5, 25), prob=0.5),
torchlm.LandmarksRandomBrightness(prob=0.),
torchlm.LandmarksRandomRotate(40, prob=0.5, bins=8),
torchlm.LandmarksRandomCenterCrop((0.5, 1.0), (0.5, 1.0), prob=0.5)
])
此外，还提供了一个用户友好的API build_default_transform来构建默认的变换流程。
transform = torchlm.build_default_transform(
input_size=(input_size, input_size),
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
force_norm_before_mean_std=True, # 首先img/=255.
rotate=30,
keep_aspect=False,
to_tensor=True # 数组 -> 张量 & HWC -> CHW
)
有关支持的变换集的更多信息，请参见transforms.md，更多示例可以在test/transforms.py中找到。
💡 关于torchlm中变换的更多详情
torchlm提供了30多种原生地标数据增强方法，并可以通过torchlm.bind方法与来自torchvision和albumentations的80多种变换进行绑定。地标的布局格式为xy，形状为(N, 2)，其中N表示输入地标的数量。此外，torchlm.bind在绑定级别提供了一个prob参数，可以强制任何变换或可调用函数成为随机风格的增强。torchlm中的数据增强是安全和最简单的。在运行时，任何导致地标超出范围的变换操作都会自动丢弃，以保持地标数量不变。是的，如果你将一个张量传递给一个类似np.ndarray的变换，torchlm会自动兼容不同的数据类型，然后通过autodtype包装器将其包装回原始类型。
绑定80多种torchvision和albumentations的变换
注意：如果你想绑定albumentations的变换，请先安装albumentations。如果你在不同版本的opencv（opencv-python和opencv-python-headless，albumentations需要opencv-python-headless）之间遇到冲突问题，请先卸载opencv-python和opencv-python-headless，然后重新安装albumentations。更多详情请参见albumentations#1140。
# 首先卸载冲突的opencv
pip uninstall opencv-python
pip uninstall opencv-python-headless
pip uninstall albumentations # 如果你已经安装了albumentations
pip install albumentations # 然后重新安装albumentations，它也会安装依赖，如opencv
然后，检查albumentations是否可用。
torchlm.albumentations_is_available() # True 或 False
transform = torchlm.LandmarksCompose([
torchlm.bind(torchvision.transforms.GaussianBlur(kernel_size=(5, 25)), prob=0.5),
torchlm.bind(albumentations.ColorJitter(p=0.5))
])
绑定自定义的可调用数组或张量变换函数
# 首先，定义你的自定义函数
def callable_array_noop(img: np.ndarray, landmarks: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: # 在这里进行一些变换 ...
return img.astype(np.uint32), landmarks.astype(np.float32)
def callable_tensor_noop(img: Tensor, landmarks: Tensor) -> Tuple[Tensor, Tensor]: # 在这里进行一些变换 ...
return img, landmarks
# 然后，绑定你的函数并将其放入变换流程中。
transform = torchlm.LandmarksCompose([
torchlm.bind(callable_array_noop, bind_type=torchlm.BindEnum.Callable_Array),
torchlm.bind(callable_tensor_noop, bind_type=torchlm.BindEnum.Callable_Tensor, prob=0.5)
])
torchlm变换的一些全局调试设置
- 全局设置日志模式为
True可能有助于你了解运行时的详细信息
# 一些全局设置
torchlm.set_transforms_debug(True)
torchlm.set_transforms_logging(True)
torchlm.set_autodtype_logging(True)
每次运行时会显示一些详细信息,信息可能如下所示:
LandmarksRandomScale() 自动数据类型信息: AutoDtypeEnum.Array_InOut
LandmarksRandomScale() 执行标志: False
BindTorchVisionTransform(GaussianBlur())() 自动数据类型信息: AutoDtypeEnum.Tensor_InOut
BindTorchVisionTransform(GaussianBlur())() 执行标志: True
BindAlbumentationsTransform(ColorJitter())() 自动数据类型信息: AutoDtypeEnum.Array_InOut
BindAlbumentationsTransform(ColorJitter())() 执行标志: True
BindTensorCallable(callable_tensor_noop())() 自动数据类型信息: AutoDtypeEnum.Tensor_InOut
BindTensorCallable(callable_tensor_noop())() 执行标志: False
LandmarksRandomTranslate() 发生错误跳过, 标志: False 错误信息: LandmarksRandomTranslate() 输入了98个关键点,但输出了96个关键点!
LandmarksRandomTranslate() 执行标志: False
- 执行标志: True 表示当前转换成功执行,False 表示由于随机概率或某些运行时异常而未执行(如果调试模式为True,torchlm 会显示错误信息)。
- 自动数据类型信息:
- Array_InOut 表示当前转换需要输入 np.ndarray 并输出 np.ndarray。
- Tensor_InOut 表示当前转换需要输入 torch Tensor 并输出 torch Tensor。
- Array_In 表示当前转换需要输入 np.ndarray 并输出 torch Tensor。
- Tensor_In 表示当前转换需要输入 torch Tensor 并输出 np.ndarray。
🎉🎉训练
在 torchlm 中,每个模型都有两个高级且用户友好的 API,分别是 apply_training 和 apply_freezing,用于训练。apply_training 处理训练过程,而 apply_freezing 决定是否冻结主干网络以进行微调。
快速开始👇
这里以 PIPNet 为例。你可以在微调之前通过 apply_freezing 冻结主干网络。
from torchlm.models import pipnet
# 如果 pretrained=True,将自动从最新版本下载预训练权重
model = pipnet(backbone="resnet18", pretrained=True, num_nb=10, num_lms=98, net_stride=32,
input_size=256, meanface_type="wflw", backbone_pretrained=True)
model.apply_freezing(backbone=True)
model.apply_training(
annotation_path="../data/WFLW/converted/train.txt", # 或微调你的自定义数据
num_epochs=10,
learning_rate=0.0001,
save_dir="./save/pipnet",
save_prefix="pipnet-wflw-resnet18",
save_interval=1,
logging_interval=1,
device="cuda",
coordinates_already_normalized=True,
batch_size=16,
num_workers=4,
shuffle=True
)
请查看 torchlm 中每个定义模型的 apply_training API 的函数入口点以获取详细文档,例如 pipnet/_impls.py#L166。如果训练过程正在运行,你可能会看到一些日志:
DataLoader 参数: {'batch_size': 16, 'num_workers': 4, 'shuffle': True}
构建 _PIPTrainDataset: 训练数量为 7500 !
第 0/9 轮
----------
[第 0/9 轮, 批次 1/468] <总损失: 0.372885> <分类损失: 0.063186> <x 损失: 0.078508> <y 损失: 0.071679> <nbx 损失: 0.086480> <nby 损失: 0.073031>
[第 0/9 轮, 批次 2/468] <总损失: 0.354169> <分类损失: 0.051672> <x 损失: 0.075350> <y 损失: 0.071229> <nbx 损失: 0.083785> <nby 损失: 0.072132>
[第 0/9 轮, 批次 3/468] <总损失: 0.367538> <分类损失: 0.056038> <x 损失: 0.078029> <y 损失: 0.076432> <nbx 损失: 0.083546> <nby 损失: 0.073492>
[第 0/9 轮, 批次 4/468] <总损失: 0.339656> <分类损失: 0.053631> <x 损失: 0.073036> <y 损失: 0.066723> <nbx 损失: 0.080007> <nby 损失: 0.066258>
[第 0/9 轮, 批次 5/468] <总损失: 0.364556> <分类损失: 0.051094> <x 损失: 0.077378> <y 损失: 0.071951> <nbx 损失: 0.086363> <nby 损失: 0.077770>
[第 0/9 轮, 批次 6/468] <总损失: 0.371356> <分类损失: 0.049117> <x 损失: 0.079237> <y 损失: 0.075729> <nbx 损失: 0.086213> <nby 损失: 0.081060>
...
[第 0/9 轮, 批次 33/468] <总损失: 0.298983> <分类损失: 0.041368> <x 损失: 0.069912> <y 损失: 0.057667> <nbx 损失: 0.072996> <nby 损失: 0.057040>
数据集格式👇
annotation_path 参数表示自定义标注文件的路径,格式必须为:
"img0_path x0 y0 x1 y1 ... xn-1,yn-1"
"img1_path x0 y0 x1 y1 ... xn-1,yn-1"
"img2_path x0 y0 x1 y1 ... xn-1,yn-1"
"img3_path x0 y0 x1 y1 ... xn-1,yn-1"
...
如果 annotation_path 中的标签已经按图像大小进行了归一化,请在 apply_training API 中将 coordinates_already_normalized 设置为 True。
"img0_path x0/w y0/h x1/w y1/h ... xn-1/w,yn-1/h"
"img1_path x0/w y0/h x1/w y1/h ... xn-1/w,yn-1/h"
"img2_path x0/w y0/h x1/w y1/h ... xn-1/w,yn-1/h"
"img3_path x0/w y0/h x1/w y1/h ... xn-1/w,yn-1/h"
...
这里有一个 WFLW 的示例,向你展示如何准备数据集,也可以参考 test/data.py。
额外的自定义设置👋
torchlm中的某些模型除了支持自定义数据集的num_lms外,还支持其他额外的自定义设置。例如,PIPNet还需要设置由自定义数据集生成的平均脸。请查看torchlm中每个定义模型的源代码,了解额外自定义设置的详细信息,以获得更灵活的训练或微调过程。以下是如何在自定义数据集上使用自定义平均脸设置来训练PIPNet的示例:
通过pipnet.set_custom_meanface方法设置自定义平均脸和最近邻关键点,该方法将计算平均脸中不同关键点之间的欧几里得距离,并自动为每个关键点设置最近邻。注意:如果自定义数据集中的关键点数量与初始化时的num_lms不相等,PIPNet将重塑检测头。
def set_custom_meanface(custom_meanface_file_or_string: str) -> bool:
"""
:param custom_meanface_file_or_string: 包含归一化或未归一化平均脸坐标的长字符串或文件,
格式为"x0,y0,x1,y1,x2,y2,...,xn-1,yn-1"。
:return: 状态,成功则为True。
"""
此外,torchlm还提供了generate_meanface API来帮助你获取自定义数据集的平均脸。
# 生成自定义平均脸
custom_meanface, custom_meanface_string = torchlm.data.annotools.generate_meanface(
annotation_path="../data/WFLW/converted/train.txt",
coordinates_already_normalized=True)
# 检查生成的平均脸
rendered_meanface = torchlm.data.annotools.draw_meanface(
meanface=custom_meanface, coordinates_already_normalized=True)
cv2.imwrite("./logs/wflw_meanface.jpg", rendered_meanface)
# 设置自定义平均脸
model.set_custom_meanface(custom_meanface_file_or_string=custom_meanface_string)
基准数据集转换器
在torchlm中,提供了一些预定义的数据集转换器,用于常用的基准数据集,如300W、COFW、WFLW和AFLW。这些转换器将帮助你将常用数据集转换为torchlm需要的标准注释格式。以下是WFLW的示例。
from torchlm.data import LandmarksWFLWConverter
# 设置原始下载数据集的路径
converter = LandmarksWFLWConverter(
data_dir="../data/WFLW", save_dir="../data/WFLW/converted",
extend=0.2, rebuild=True, target_size=256, keep_aspect=False,
force_normalize=True, force_absolute_path=True
)
converter.convert()
converter.show(count=30) # 显示一些转换后的图像及关键点以进行调试
然后,../data/WFLW/converted中的输出布局将如下所示:
├── image
│ ├── test
│ └── train
├── show
│ ├── 16--Award_Ceremony_16_Award_Ceremony_Awards_Ceremony_16_589x456y91.jpg
│ ├── 20--Family_Group_20_Family_Group_Family_Group_20_118x458y58.jpg
...
├── test.txt
└── train.txt
推理
C++ API
torchlm的ONNXRuntime(CPU/GPU)、MNN、NCNN和TNN C++推理将在**lite.ai.toolkit中发布。以下是使用FaceLandmarks1000进行1000个面部关键点检测**的示例。从Model-Zoo下载模型2。
#include "lite/lite.h"
static void test_default()
{
std::string onnx_path = "../../../hub/onnx/cv/FaceLandmark1000.onnx";
std::string test_img_path = "../../../examples/lite/resources/test_lite_face_landmarks_0.png";
std::string save_img_path = "../../../logs/test_lite_face_landmarks_1000.jpg";
auto *face_landmarks_1000 = new lite::cv::face::align::FaceLandmark1000(onnx_path);
lite::types::Landmarks landmarks;
cv::Mat img_bgr = cv::imread(test_img_path);
face_landmarks_1000->detect(img_bgr, landmarks);
lite::utils::draw_landmarks_inplace(img_bgr, landmarks);
cv::imwrite(save_img_path, img_bgr);
delete face_landmarks_1000;
}
更多人脸对齐类（68点、98点、106点、1000点）
auto *align = new lite::cv::face::align::PFLD(onnx_path); // 106个关键点，仅1.0Mb！
auto *align = new lite::cv::face::align::PFLD98(onnx_path); // 98个关键点，仅4.8Mb！
auto *align = new lite::cv::face::align::PFLD68(onnx_path); // 68个关键点，仅2.8Mb！
auto *align = new lite::cv::face::align::MobileNetV268(onnx_path); // 68个关键点，仅9.4Mb！
auto *align = new lite::cv::face::align::MobileNetV2SE68(onnx_path); // 68个关键点，仅11Mb！
auto *align = new lite::cv::face::align::FaceLandmark1000(onnx_path); // 1000个关键点，仅2.0Mb！
auto *align = new lite::cv::face::align::PIPNet98(onnx_path); // 98个关键点，CVPR2021！
auto *align = new lite::cv::face::align::PIPNet68(onnx_path); // 68个关键点，CVPR2021！
auto *align = new lite::cv::face::align::PIPNet29(onnx_path); // 29个关键点，CVPR2021！
auto *align = new lite::cv::face::align::PIPNet19(onnx_path); // 19个关键点，CVPR2021！
更多C++ API详情，请查看 lite.ai.toolkit。 
Python API👇
在 torchlm 中，我们提供了使用 PyTorch 和 ONNXRuntime 部署模型的流程。一个名为 runtime.bind 的高级API可以将人脸检测和关键点模型绑定在一起，然后你可以运行 runtime.forward API来获取输出的关键点和边界框。这里是 PIPNet 的一个示例。PIPNet的预训练权重，下载。
在PyTorch后端上推理
import torchlm
from torchlm.tools import faceboxesv2
from torchlm.models import pipnet
torchlm.runtime.bind(faceboxesv2(device="cpu")) # 如果你想用CUDA运行，设置device="cuda"
# 如果你想用CUDA运行，设置map_location="cuda"
torchlm.runtime.bind(
pipnet(backbone="resnet18", pretrained=True,
num_nb=10, num_lms=98, net_stride=32, input_size=256,
meanface_type="wflw", map_location="cpu", checkpoint=None)
) # 如果pretrained=True，将自动从最新版本下载预训练权重
landmarks, bboxes = torchlm.runtime.forward(image)
image = torchlm.utils.draw_bboxes(image, bboxes=bboxes)
image = torchlm.utils.draw_landmarks(image, landmarks=landmarks)
在ONNXRuntime后端上推理
import torchlm
from torchlm.runtime import faceboxesv2_ort, pipnet_ort
torchlm.runtime.bind(faceboxesv2_ort())
torchlm.runtime.bind(
pipnet_ort(onnx_path="pipnet_resnet18.onnx",num_nb=10,
num_lms=98, net_stride=32,input_size=256, meanface_type="wflw")
)
landmarks, bboxes = torchlm.runtime.forward(image)
image = torchlm.utils.draw_bboxes(image, bboxes=bboxes)
image = torchlm.utils.draw_landmarks(image, landmarks=landmarks)
🤠🎯 评估
在 torchlm 中，每个模型都有一个名为 apply_evaluating 的高级且用户友好的API用于评估。这个方法将计算评估数据集的NME、FR和AUC。这里是 PIPNet 的一个示例。
from torchlm.models import pipnet
# 如果pretrained=True，将自动从最新版本下载预训练权重
model = pipnet(backbone="resnet18", pretrained=True, num_nb=10, num_lms=98, net_stride=32,
input_size=256, meanface_type="wflw", backbone_pretrained=True)
NME, FR, AUC = model.apply_evaluating(
annotation_path="../data/WFLW/convertd/test.txt",
norm_indices=[60, 72], # 两个眼球的索引。
coordinates_already_normalized=True,
eval_normalized_coordinates=False
)
print(f"NME: {NME}, FR: {FR}, AUC: {AUC}")
然后，你将得到 性能(@NME@FR@AUC) 结果。
Built _PIPEvalDataset: eval count is 2500 !
Evaluating PIPNet: 100%|██████████| 2500/2500 [02:53<00:00, 14.45it/s]
NME: 0.04453323229181989, FR: 0.04200000000000004, AUC: 0.5732673333333334
⚙️⚔️ 导出
在 torchlm 中，每个模型都有一个名为 apply_exporting 的高级且用户友好的API用于ONNX导出。这里是 PIPNet 的一个示例。
from torchlm.models import pipnet
# 如果pretrained=True，将自动从最新发布版本下载预训练权重
model = pipnet(backbone="resnet18", pretrained=True, num_nb=10, num_lms=98, net_stride=32,
input_size=256, meanface_type="wflw", backbone_pretrained=True)
model.apply_exporting(
onnx_path="./save/pipnet/pipnet_resnet18.onnx",
opset=12, simplify=True, output_names=None # 使用默认输出名称
)
然后，如果导出过程完成，你将得到一个静态ONNX模型文件。
...
%195 = Add(%259, %189)
%196 = Relu(%195)
%outputs_cls = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [1, 1]](%196, %cls_layer.weight, %cls_layer.bias)
%outputs_x = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [1, 1]](%196, %x_layer.weight, %x_layer.bias)
%outputs_y = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [1, 1]](%196, %y_layer.weight, %y_layer.bias)
%outputs_nb_x = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [1, 1]](%196, %nb_x_layer.weight, %nb_x_layer.bias)
%outputs_nb_y = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [1, 1]](%196, %nb_y_layer.weight, %nb_y_layer.bias)
return %outputs_cls, %outputs_x, %outputs_y, %outputs_nb_x, %outputs_nb_y
}
检查0/3...
检查1/3...
检查2/3...
📖 文档
🎓 许可证
torchlm的代码以MIT许可证发布。
❤️ 贡献
如果你喜欢这个项目，请考虑给它⭐，这是支持我的最简单方式。
👋 致谢
- torchlm的transforms实现借鉴了Paperspace的代码。
- PIPNet：Towards Efficient Facial Landmark Detection in the Wild, CVPR2021
