如何利用树莓派创建自定义关键点检测AI模型

欢迎来到本教程，在这里您将学习如何在树莓派AI相机上创建和部署您自己的自定义关键点检测AI模型。

构建在树莓派AI相机等边缘设备上高效运行的AI模型可能具有挑战性。为了简化这个过程，我们开发了样例代码和工具来优化整个工作流程——从培训到部署。

本文是我们关于树莓派AI相机的AI模型开发系列的第2部分。第1部分介绍了使用Nanodet在树莓派AI相机上创建和运行自定义对象检测AI模型。

本文概述

•使用我们的示例代码和分步说明创建自定义关键点检测模型

•生成优化的模型文件，准备部署在树莓派AI相机

•在设备上实现并运行您训练过的模型

理解关键点检测AI模型

关键点人工智能模型是一种人工智能模型，它可以检测到人类或物体上不同的兴趣点(关键点)，并将它们的位置作为坐标输出。

树莓派AI相机配备了专门为人体姿势估计设计的即用型预训练模型。这些模型可以检测骨骼关键点，包括头、肩膀、肘部、手腕、臀部、膝盖和脚踝。

这些模型是实时工作的，非常适合姿势分析、手势识别和运动跟踪等应用。

使用预训练模型

如果你想在不训练自己的模型的情况下检测人体姿势，你可以使用树莓派提供的预构建模型。这些部队无需额外训练即可立即部署。

通过再培训扩大应用范围

虽然预先训练的模型在人体姿势检测方面表现出色，但如果你想检测其他物体上的关键点，你需要训练自己的模型。

通过使用自己的数据集进行再训练，您可以检测几乎任何对象或主题上的关键点。这带来了许多可能性：

•读取生产线上的模拟仪表和表盘

•估计机器人操作和控制的关节位置

•跟踪宠物或野生动物的姿势和运动模式

•检测制造零件的对中点以进行质量控制

您可以定义对应用程序重要的关键点，使关键点检测适用于任何需要对特定形状进行精确点检测的应用程序。

在本教程中，我们将从一个基本的例子开始：检测箭头形状的角点。这个简单实用的用例将帮助您在移动到更复杂的应用程序之前了解边缘AI关键点检测的基础知识。

从示例数据开始

我们提供了一个现成的箭头数据集，以便您可以快速测试工作流。一旦理解了这个过程，就可以轻松地将自己的数据集替换为任何自定义的关键点检测任务。

系统需求

为部署

•树莓派(与AI相机兼容的任何型号)

•树莓派AI相机

用于模型训练

•带有NVIDIA GPU的计算机(强烈推荐用于高效训练)

•Ubuntu 22.04(或兼容的Linux发行版)

•Python 3.10

注意：虽然GPU显著加快训练速度，但你可以在CPU上训练，但可能需要更长的时间。

本地环境建设

注意：本节介绍在您的机器上本地设置培训环境。如果你已经在我们的自定义对象检测AI模型教程中设置了本地环境，那么你可以跳过本节。

1. 克隆存储库

2. 设置

安装必要的软件包。

3. 创建并激活Python 3.10虚拟环境

本教程假设使用Python 3.10，所以首先，让我们确认已经安装了3.10。

注意：由于所需库的版本依赖关系，请确保使用Python 3.10。

如果显示如下，则表示安装了Python 3.10：

如果你还没有安装Python 3.10，你可以按照以下步骤安装，例如：

接下来，创建一个虚拟环境。

4. 安装软件包

步骤1：训练模型

示例存储库在许多.ini配置文件中提供了与AI模型训练、量化和评估相关的各种设置。

在本教程中，我们将使用posenet_arrow.ini来训练模型。

要点

其特点是，您可以通过更改设置灵活地调整训练条件，而无需直接重写Python代码。您可以选择数据集和任务，并调整参数。在使用自己的数据集进行训练一节中解释了这种自定义方法。

使用数据集

用于此训练的数据集是箭头形状的数据集。

你可以去这里获取原始图像。

在样本数据集中，每个顶点被分配了7个点，如下图所示。

移动到存储ini文件的samples文件夹，运行以下命令读取ini文件并开始训练。

下面的命令将根据指定的.ini文件对模型进行训练和量化。

当运行时，这将下载数据集并开始训练。默认设置文件指定50个epoch。

这会生成以下模型文件：

步骤2：量化和转换训练模型

将量化模型转换并打包成与树莓派AI相机兼容的格式。

注意：继续在venv环境中工作，使用与用于训练的相同的TensorFlow版本进行量化。版本不匹配导致错误。

本教程解释了将a.keras转换为网络的过程。rpk，我们可以上传到IMX500。

1. 安装Edge-MDT(模型开发工具包)

安装Edge-MDT(模型开发工具包)，包括量化、压缩和转换AI模型所需的工具。PyTorch和TensorFlow都有可用的版本，我们将安装TensorFlow版本：

2. 执行模型量化

下面的命令执行模型量化，-o标志指定创建和存储AI模型文件的目录。在下面的示例中，它们将被输出到convert_result。

转换大约需要30秒。

这会生成以下文件：

当packerOut.zip生成后，步骤2完成。

步骤3:AI相机包装

注意：从这里开始，操作将在Raspberry Pi⭐️上进行

将转换后的文件夹移动到树莓派。在本教程中，文件夹名称为convert_result。

将步骤2中生成的packerOut.zip文件移动到树莓派上的任何位置。您可以使用scp或任何您喜欢的工具进行文件传输。

如果您没有安装，请在树莓派上安装模型打包工具：

执行转换，将模型打包到RPK文件中。下面的命令指定packerOut.zip作为输入，并将打包的输出文件保存到rpk_output_folder中。

生成一个打包的AI模型。rpk扩展。

当网络。生成rpk文件，则完成步骤3。

步骤4：在树莓派上执行

让我们在Raspberry Pi上安装必要的库。

为了可视化关键点，我们将使用Modlib，一个用于树莓派AI相机应用程序开发的库。该库附带了几个样本检测模型，并允许从预训练模型到自定义模型轻松实现。在本教程中，我们将使用这个库来简单实现AI模型可视化。

1. 安装Modlib:

2. 准备一个可视化脚本

基于下面的示例脚本创建一个文件visualize_arrow.py。

这个脚本检测箭头的七个顶点，并用不同的颜色将它们可视化。您需要在您自己的脚本中更改rpk AI模型文件的路径。

示例可视化脚本

3. 执行visualize_arrow.py:

结果

当您运行脚本时，AI模型自动部署到您的树莓派AI相机，并出现一个预览窗口，显示实时检测。

这是你将看到的：

关键点被准确地检测到每个顶点，确认您的自定义训练模型工作正常。

用你自己的数据集进行训练

准备好使用自己的数据集了吗?这是需要改变的：

•更新.ini配置文件：

•[DATASET] NAME = YourDatasetName

•[MODEL] CLASS_NUM =类的个数

•[TRAINER] CONFIG =你的YAML配置文件

结论

在本教程中，您已经学习了如何使用我们的再培训工具和示例代码在树莓派AI相机上创建和部署自定义关键点检测模型。

虽然我们将箭头检测作为学习练习，但您学到的技术可以应用于无数的现实场景。只需交换数据集，您就可以为无数边缘AI应用程序创建关键点检测解决方案。

本文编译自hackster.io