Press "Enter" to skip to content

智能切图平台 – 基于yolov5神经网络的目标检测实践总结

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

在线切图平台项目是前端通道的提效项目,希望能通过对设计稿的分析,将设计稿转成HTML,我们尝试使用深度学习来自动识别图片中的文字和图片,这篇文章记录实现过程和总结。

 

yolov5是什幺

 

YOLO 系列是深度学习目标检测领域里比较出众算法系列,最初是 Facebook 的一位大牛的作品,目前总共有个五个版本,对它的具体了解可以参考下面的文章:

 

写给小白的YOLO介绍

 

简单说就是一个可以检测出图片中有哪些物体的算法模型,而且效率很高,可以实时视频检测。

 

参考下面的视频直观体验:

 

YOLOv5视频体验

 

样本准备

 

我们准备了 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 的监控图:

训练结果

 

训练结束后,可以看到训练的结果数据:

这里面主要看 PrecisionRecall 这两个指标:

 

Precision
Recall

 

本次训练 2000 轮,准确率能达到 68%,也就是检测结果中,有 68% 是正确的,召回率 60% ,准确的结果占总样本 60%。

 

这个准确率是还是有提升空间,达到 90% 以上是我们的目标。

 

验证

 

验证命令:

 

python detect.py --source ./inference/images/ --weights yolov5s.pt --conf 0.4

source 设置图片的来源
weights 设置训练生成的权重文件
conf 设置检查的阈值,高于这个阈值的都显示出来

在测试集合上跑一下,贴几个检查结果:

从结果来看,大部分结果是符合预期的,效果还不错。

 

后续规划

 

后续主要的目标提供准确率,增加一些样本数量,提高训练的轮数,然后将检测程序部署到线上提供服务。

 

虽然依靠深度学习可以检测出图片的目标区域,但是这个不是精确的,框选的范围和实际范围是有误差的,后续会通过其他的手段校准这个框。

 

总结

 

本文大致梳理了最近采用 yolov5 实现图片目标检测的过程,过程还是比较曲折的,毕竟对深度学习不太了解,前期要补课,了解神经网络相关的基础原理和概念,后期才动手操作,关于深度学习,目前业界已经有比较成熟的平台了,如果想做工程,可是动手试试,简单数字识别、物体识别、图像分类等任务都可以做到,对算法深度学习和优化,这个应该是需要算法专家来做了,还好公司大牛可以咨询,每次有问题都能给指个道。

 

资源推荐

 

深度学习资源:

 

视频(吴恩达的两门课):

https://www. bilibili.com/video/BV16 4411b7dx?p=1
https://www. bilibili.com/video/BV16 4411m79z?p=1

Be First to Comment

发表评论

邮箱地址不会被公开。 必填项已用*标注