torchlm

torchlm-logo

<div align='center'> <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg > <img src=https://img.shields.io/pypi/v/torchlm?color=aff > <img src=https://static.pepy.tech/personalized-badge/torchlm?period=total&units=international_system&left_color=grey&right_color=blue&left_text=Downloads > <img src=https://img.shields.io/pypi/pyversions/torchlm?color=dfd > <img src=https://img.shields.io/badge/macos|linux|windows-pass-skyblue.svg > <img src=https://img.shields.io/badge/license-MIT-lightblue.svg > </div> <p align="center">英文 | <a href="docs/api/transforms.md">数据增强API文档</a> | <a href="https://www.zhihu.com/column/c_1426666301352218624">知乎专栏</a> | <a href="https://pepy.tech/project/torchlm">Pypi下载量</a> </p>

🤗 简介

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中找到。

<details> <summary>💡 关于torchlm中变换的更多详情 </summary>

torchlm提供了30多种原生地标数据增强方法，并可以通过torchlm.bind方法与来自torchvision和albumentations的80多种变换进行绑定。地标的布局格式为xy，形状为(N, 2)，其中N表示输入地标的数量。此外，torchlm.bind在绑定级别提供了一个prob参数，可以强制任何变换或可调用函数成为随机风格的增强。torchlm中的数据增强是安全和最简单的。在运行时，任何导致地标超出范围的变换操作都会自动丢弃，以保持地标数量不变。是的，如果你将一个张量传递给一个类似np.ndarray的变换，torchlm会自动兼容不同的数据类型，然后通过autodtype包装器将其包装回原始类型。

<details> <summary> 绑定80多种torchvision和albumentations的变换 </summary>

注意：如果你想绑定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))
])

</details> <details> <summary> 绑定自定义的可调用数组或张量变换函数 </summary>

# 首先，定义你的自定义函数
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)
])

</details> <details> <summary> torchlm变换的一些全局调试设置 </summary>

全局设置日志模式为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。
是的,你可以将 Tensor 传递给类似 np.ndarray 的转换,torchlm 会通过 autodtype 封装器自动兼容不同的数据类型,并将其包装回原始类型。

🎉🎉训练

在 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下载模型<sup>2</sup>。

#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...