智能切图平台 – 基于yolov5神经网络的目标检测实践总结
在线切图平台项目是前端通道的提效项目，希望能通过对设计稿的分析，将设计稿转成HTML，我们尝试使用深度学习来自动识别图片中的文字和图片，这篇文章记录实现过程和总结。
yolov5是什么
YOLO 系列是深度学习目标检测领域里比较出众算法系列，最初是 Facebook 的一位大牛的作品，目前总共有个五个版本，对它的具体了解可以参考下面的文章：
简单说就是一个可以检测出图片中有哪些物体的算法模型，而且效率很高，可以实时视频检测。
参考下面的视频直观体验：
样本准备
我们准备了 1000 个线上站点的截图作为样本，利用了 apify 这个爬虫库，对每个网页截屏， apify 是基于 puppeteer 实现的，有兴趣可以研究下，贴段代码。
const Apify = require('apify');
Apify.main(async () => {
const requestQueue = await Apify.openRequestQueue();
await requestQueue.addRequest({ url: ' http://you/site' });
const pseudoUrls = [
new Apify.PseudoUrl('http://allow.site.com/[.*]'),
];
let i = 0;
const crawler = new Apify.PuppeteerCrawler({
requestQueue,
handlePageFunction: async ({ request, response, page }) => {
const status = response.status();
const title = await page.title();
console.log(`Title of ${request.url}: ${title}`);
await page.setViewport({
width: 1296,
height: 300,
deviceScaleFactor: 0.5,
})
await page.screenshot({ path: `../images/img-${i++}.jpg`, fullPage: true, quality: 100 });
await Apify.utils.enqueueLinks({ page, selector: 'a', pseudoUrls, requestQueue });
page.close();
},
maxRequestsPerCrawl: 999999,
maxConcurrency: 10,
gotoTimeoutSecs: 60
});
await crawler.run();
});
一千个样本截图：
样本标注
使用labelimg软件对样本进行标注，labelimg 是基于 Python+qt 实现的，项目地址如下：
https://github.com/tzutalin/labelImg
windows 上可以参考 Windows + Anaconda 的安装方式。
图像界面如下：
样本标注会生成 label 文件，标识出每个框的类别，这里我们先确定了五个类别，包括：image,text,container,list,table
CPU上跑起来
在 WPAI 平台上训练之前，先尝试在本地跑起来，尝试用 100 个样本进行训练，跑这个训练最好是用 linux 系统，因为最后都是要跑在linux 系统上，我用的 win10 自带的 linux 子系统，如果 linux 子系统安装不上，可以用虚拟机 VMWare 或者 VirtualBox 安装 Ubuntu 或者 CentOS。
怎么跑训练数据，在 yolov5 的github 上有讲解：
https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
训练命令：
python train.py --img 640 --batch 5 --epochs 100 --data ./data/pages.yaml --cfg ./models/yolov5s.yaml --weights ./weights/best.pt --worker 1
我在本地训练时遇到几个问题：
maximum number of dataloader workers
在 CPU 上训练，效率很低，100 个样本跑 2000 轮，大概需要四天左右，基本上能看到训练的结果，比如：
看到这个识别效果，很是兴奋啊！！
WPAI平台离线训练
首先得先感谢 WAPAI 和 凤凰平台的团队，为我们提供了一个这样好的环境。
在本地测试没问题后，考虑在 WAPAI 上开展训练，WPAI 有 GPU 环境，希望能感受一下 GPU 的效率。
首先，建立调试环境，调试环境选择 PyTorch1.6 容器，将训练需要的程序、样本上传到调试环境，调试环境主要是进行一些测试，提供了 jupyter 环境，可以在线写 Python 脚本， jupyter 也提供了命令行工具，将 yolov5 的代码以及准备的好样本打包，利用 jupter 上传到对应的目录下。利用命令行解压，然后在调试环境跑一下试试。
然后，建立离线环境，离线环境提供多个平台，我们选择的是 PyTorch 单机版，下面是配置表单：
资源配置：
- GPU 资源，可以选择 CPU 或者 GPU，GPU 可以选择 K40 和 P40，我们选择 P40，因为PyTorch自从1.3以后，不再支持K40了
- CPU梳理，选择 2 个
- 内存 5 G
- 镜像环境选择 PyTorch1.6 GPU版
模型训练配置：
- 训练方式：选择执行调试环境程序
- 工作目录：这个目录是调试环境内设置的目录，是在 jupyter 中看到目录，不是系统中目录的绝对路径，jupyter 的的工作目录默认是 /code ，但是这里不能设置为/code/yolv5
- 训练程序：是相对工作目录的文件路径
- 程序参数：这里是传给训练程序的参数，batch 设置为16，16 张图作为一个批次，worker位置为2，logdir 单独指定一个目录，作为日志和权重文件输出目录
- 模型路径：这个模型路径是模型生成的路径，后续训练结束会将这个目录打包提供下载，设置为和logdir一个目录。
其他配置： – tensorboard配置,tensorboard提供了对训练过程的可视化监控，将 tensorboard logdir 设置为前面的 logdir 即可。
设置好以后，就可以启动了，启动之初会安装依赖的包，安装完成后，启动训练。
每次运行，都会新建一个容器环境，所以要重新安装依赖的包，这里或许可以优化一下。
使用 GPU 训练相当的快，之前一个批次 5 张图，需要 15 秒，现在一个批次 16 张图，只需要 2秒。训练 2000 轮不到一天，
下面是 tensorboard 的监控图：
训练结果
训练结束后，可以看到训练的结果数据：
这里面主要看 Precision 和 Recall 这两个指标：
Precision Recall
本次训练 2000 轮，准确率能达到 68%，也就是检测结果中，有 68% 是正确的，召回率 60% ，准确的结果占总样本 60%。
这个准确率是还是有提升空间，达到 90% 以上是我们的目标。
验证
验证命令：
python detect.py --source ./inference/images/ --weights yolov5s.pt --conf 0.4
- source 设置图片的来源
- weights 设置训练生成的权重文件
- conf 设置检查的阈值，高于这个阈值的都显示出来
在测试集合上跑一下，贴几个检查结果：
从结果来看，大部分结果是符合预期的，效果还不错。
后续规划
后续主要的目标提供准确率，增加一些样本数量，提高训练的轮数，然后将检测程序部署到线上提供服务。
虽然依靠深度学习可以检测出图片的目标区域，但是这个不是精确的，框选的范围和实际范围是有误差的，后续会通过其他的手段校准这个框。
总结
本文大致梳理了最近采用 yolov5 实现图片目标检测的过程，过程还是比较曲折的，毕竟对深度学习不太了解，前期要补课，了解神经网络相关的基础原理和概念，后期才动手操作，关于深度学习，目前业界已经有比较成熟的平台了，如果想做工程，可是动手试试，简单数字识别、物体识别、图像分类等任务都可以做到，对算法深度学习和优化，这个应该是需要算法专家来做了，还好公司大牛可以咨询，每次有问题都能给指个道。
资源推荐
深度学习资源：
视频（吴恩达的两门课）：
书：
About The Author
