RK3399pro实现安全帽识别

hq1978

• RK3399pro实现安全帽识别
  

前段时间公司智慧工地项目要实现安全帽识别，使用 firefly 的RK3399pro的边缘盒子。最初我们找算法公司买现成的算法，由于我们的使用场景与环境非常复杂，所以实际使用时总是有问题。算法公司不愿意免费优化算法，所以公司决定自己解决。

我们公司没有搞人工智能算法的，按着网上找的什么tensorflow、pytorch 、yolo 我们小组越学越糊涂。说实话我们真心想学会，算法工程师工资高，学会了怎么训练算法并移植到边缘盒子，这样可以找一个工资更高的工作。但臣妾真的学不会呀 😦

后来找到了这个工具非常容易地训练出算法并集成到我们的系统中了。工具名我就不写了，以免广告嫌疑，直接上下载地址： https://pan.baidu.com/s/1ZPcGjUUMUQa2ndxb4TUuHQ?pwd=8888 。里面的子目录“起点AI Builder 社区版”是安装文件；子目录“训练数据”是标注好的一些常用数据，用这些数据就可以直接训练一些常用算法，真是业界良心。

安装

​安装很简单，点击下载setup.exe，然后下一步、下一步就可以安装了。
安装后有会有两个程序：
AI Builder : 训练算法用的。
图片标注工具：标注训练数据用的。

有一点要注意，如果自己的电脑安装有英伟达芯片的显卡就勾选CUDA，其他芯片的显卡就不要勾选。我的电脑电脑是RTX1070 8G显存。

我另外一个同事电脑i3的CPU，没有独立显卡，也能很好的训练出算法。
以前我们被算法公司忽悠，我一直以为训练算法的机器至少要 i7已上的CPU，要好几块RTX3090的显卡才行。现在才知道普通PC就可以训练AI算法了。

训练算法

这个工具完全颠覆了我对算法的认知，感觉是个人就能训练出AI算法。
最多就三步：标注数据、训练算法、导出

第一步：标注数据

收集足够的图片放到一个目录中，打开“图片标注工具”标注这些数据。只要会用鼠标拉框就能标注数据，很简单。
网盘中有标注好的安全帽数据，我就直接拿来用了。标注数据也省了。
这个数据中把没戴安全帽的头标注为 head, 把戴了安全帽的头标注为 helmet, 没戴在头上的安全帽没有标注。这样算法就不会把没戴在头上的安全帽识别为戴了安全帽的头。

第二步：训练算法

• 新建训练任务
  输入任务名称，我起名“安全帽识别”。
  
  选择算法要运行的设备，我用的是3399pro边缘盒子，选择“RK3399-PRO NPU”。
  如果要训练能在手机上运行的算法，可以选择“Android 手机”。这个工具真的很良心，还能做手机上的算法。
  
  选择模型大小。
  
  创建好训练任务后，将训练数据设置为收集的安全帽图片的目录。
  
• 点击训练按钮
  点击“开启训练”按钮就开始训练了。这太简单了吧 😃
  
  

第三步：导出模型与SDK

训练差不多的时候就可以导出模型。
点击“导出模型/SDK”按钮

在弹出的对话框中，设置要导出的内容与要保存的目录。
我们的项目使用的是C++语言，所以我选择“C++ SDK”。
设置好后，点击“导出”按钮。

说明文档、SDK、模型、Demo源码都会保存到指定的目录中。

集成算法

导出的文件 sdk.zip 里面有如下文件

• 文件 model.model：训练好的算法模型文件。
• class_list.txt：模型识别的物体的类别列表。该文件列出了类别编号与类别名的对应关系。
• 目录 libs ： 起点物体识别SDK的库文件与用到的第三方库。
• 目录 demo：是使用 SDK 的范例代码。
• 文件 get_key ： 用于获取设备 licnese key 的程序。
  

调用流程



范例代码

下面是使用训练出来的模型的范例代码：

1. #include <sys/time.h>
   
2. #include <iostream>
   
3. #include <fstream>
   
4. #include <vector>
   
5. #include <string>
   
6. #include <chrono>
   
7. #include "opencv2/opencv.hpp"
   
8. #include "qdian/detector.h"
   
9. 
   
10. #include <stdio.h>
    
11. #include <dirent.h>
    
12. 
    
13. std::vector<std::string> class_names;
    
14. void detect_and_draw(const char *image_file_name, const char *out_file_name);
    
15. 
    
16. int main(int argc, char **argv)
    
17. {
    
18. 
    
19.     //获取license key
    
20.     char *key;
    
21.     key = get_key();
    
22.     printf("授权key: %s\n", key);
    
23. 
    
24.         int ret;
    
25.     // 初始化AI 引擎
    
26.     // "model.model"：训练出来的模型文件
    
27.     // "lic.lic": 授权文件，把上面的key给到客服就可以获得授权文件
    
28.     // 0.55: 置信度阈值。官方给的demo 是0.35，我测试的用0.55-0.8之间会比较好
    
29.     // 0.70：重合度阈值。
    
30.     ret = init_detector("model.model", "lic.lic", 0.55, 0.70);
    
31.     if (ret == 1)
    
32.     {
    
33.         printf("init_detector OK\n");
    
34.     }
    
35.     else if (ret == 0)
    
36.     {
    
37.         printf("init_detector failed\n");
    
38.         return -1;
    
39.     }
    
40.     else if (ret == -1)
    
41.     {
    
42. 
    
43.         printf("license failed\n");
    
44.         return -1;
    
45.     }
    
46. 
    
47.     //获取模型的能识别的物体的类别名
    
48.     class_names = get_class_list_of_model("model.model");
    
49.         // 加载图片
    
50.     cv::Mat frame;
    
51.     frame = cv::imread("test.jpg", 1);
    
52. 
    
53.         // 物体识别
    
54.     std::vector<BoxInfo> result = detect(frame);
    
55.     draw_objects(result, frame, "result.jpg");
    
56. 
    
57.     // 销毁AI 引擎
    
58.     destroy_detector();
    
59. 
    
60.     return 0;
    
61. }
    
62. 
    
63. void draw_objects(std::vector<BoxInfo> result, cv::Mat frame, const char *out_file_name)
    
64. {
    
65.     //将识别的物体画在图上
    
66.     for (int i = 0; i < result.size(); i++)
    
67.     {
    
68.         // 识别出来的物体的位置
    
69.         BoxInfo obj = result[i];
    
70.         printf("%d = %.5f at %.2f %.2f %.2f  %.2f\n", obj.label, obj.score,
    
71.                obj.x1, obj.y1, obj.x2, obj.y2);
    
72.         cv::rectangle(frame, cv::Point(obj.x1, obj.y1), cv::Point(obj.x2, obj.y2), cv::Scalar(255, 0, 0));
    
73. 
    
74.         // 获取物体的类别
    
75.         char text[256];
    
76.         sprintf(text, "%s %.1f%%", class_names[obj.label].c_str(), obj.score * 100);
    
77. 
    
78.         int baseLine = 0;
    
79.         cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
    
80. 
    
81.         // 标注物体区域
    
82.         int x = obj.x1;
    
83.         int y = obj.y1 - label_size.height - baseLine;
    
84.         if (y < 0)
    
85.             y = 0;
    
86.         if (x + label_size.width > frame.cols)
    
87.             x = frame.cols - label_size.width;
    
88.         cv::rectangle(frame, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
    
89.                       cv::Scalar(255, 255, 255), -1);
    
90. 
    
91.         // 标注物体类别
    
92.         // 备注：类别名为中文会显示乱码。
    
93.         // 如果要正常显示中文，可使用 freetype 库
    
94.         cv::putText(frame, text, cv::Point(x, y + label_size.height),
    
95.                     cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
    
96.     }
    
97. 
    
98.     cv::imwrite(out_file_name, frame);
    
99. }

复制代码

模型不准怎么办

用这个工具最大的好处是，自己可以随时针对真实使用场景提升模型的精度。这样项目的落地就有了很好的保障。

方法一

最简单的办法就是将识别有问题的图片找出来，放到训练用的数据目录中，并用配套的“图片标注工具”标注好，然后重新训练。（如果是视频，就将视频解码为图片，或者直接截图）

方法二

新建一个目录，将真实使用的摄像头的视频解码为图片（大概500张）。将训练数据目录设置为改目录，然后重新导出模型。不用标注数据也不用重新训练模型，只要重新导出就可以了。
这个方法我开始不清楚是什么原因，后来问了作者，才知道用 RK3399pro NPU要把模型进行 int8 量化，量化时用的图片与真正使用时的越一样就越准。什么是模型int8量化，自己去网上查。

意外惊喜

如果导出 restFul SDK，集成起来更加简单，几乎直接可以用，细节大家可以仔细研究。这个工具越用越觉得好。

其他问题

360 会误报误杀。安装与使用时建议关闭360。
这事情我专门去问了作者，作者说是里面配套的数据标注软件是python的，为了方便使用，将python代码打包为了exe文件，所以360会误报误杀。

感想

在以前以为AI算法高深莫测，所有的算法公司与算法工程师都告诉我们训练AI算法需要多高的算力，要学习一堆高深的数学知识，我们这个小公司根本没办法自己单独训练算法的。我们也一直认为我们无法在没有AI算法工程师的情况下训练算法的。直到老板的朋友推荐这个工具给我们后，我们才知道原来自己定制AI算法原来可以这么简单。

wx__JDZXBp

不错，谢谢分享

skyzhao

感谢分享

Ruan1988

先收藏，有需要再拿