使用YOLOV5进行垃圾满溢检测

最近老板让做一个社区中垃圾满溢检测的demo，但因为没有数据，本身的研究方向也不是目标检测，因此在学习过程中写了这篇博客，记录demo整个流程。

1. 数据收集

1.1 数据来源

由于没有找到公开的数据集用于训练，因此使用在百度中爬取的图片重新进行标注。

1.2 以“垃圾桶”为关键字在百度中爬取图片

首先看有没有别人写好的代码可以直接用，之后在知乎上找到一篇文章详细的介绍了怎么爬取百度的图片，并且给出了代码，试用之后发现完全够解决需求，在这里给出代码，具体的解释可以查看原文爬虫代码原文。

import requests
import os
import re


def get_images_from_baidu(keyword, page_num, save_dir):
    # UA 伪装：当前爬取信息伪装成浏览器
    # 将 User-Agent 封装到一个字典中
    # 【（网页右键 → 审查元素）或者 F12】 → 【Network】 → 【Ctrl+R】 → 左边选一项，右边在 【Response Hearders】 里查找
    header = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/78.0.3904.108 Safari/537.36'}
    # 请求的 url
    url = 'https://image.baidu.com/search/acjson?'
    n = 200000
    for pn in range(0, 30 * page_num, 30):
        # 请求参数
        param = {'tn': 'resultjson_com',
                 # 'logid': '7603311155072595725',
                 'ipn': 'rj',
                 'ct': 201326592,
                 'is': '',
                 'fp': 'result',
                 'queryWord': keyword,
                 'cl': 2,
                 'lm': -1,
                 'ie': 'utf-8',
                 'oe': 'utf-8',
                 'adpicid': '',
                 'st': -1,
                 'z': '',
                 'ic': '',
                 'hd': '',
                 'latest': '',
                 'copyright': '',
                 'word': keyword,
                 's': '',
                 'se': '',
                 'tab': '',
                 'width': '',
                 'height': '',
                 'face': 0,
                 'istype': 2,
                 'qc': '',
                 'nc': '1',
                 'fr': '',
                 'expermode': '',
                 'force': '',
                 'cg': '',    # 这个参数没公开，但是不可少
                 'pn': pn,    # 显示：30-60-90
                 'rn': '30',  # 每页显示 30 条
                 'gsm': '1e',
                 '1618827096642': ''
                 }
        request = requests.get(url=url, headers=header, params=param)
        if request.status_code == 200:
            print('Request success.')
        request.encoding = 'utf-8'
        # 正则方式提取图片链接
        html = request.text
        image_url_list = re.findall('"thumbURL":"(.*?)",', html, re.S)
        print(image_url_list)
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        for image_url in image_url_list:
            image_data = requests.get(url=image_url, headers=header).content
            with open(os.path.join(save_dir, f'{n:06d}.jpg'), 'wb') as fp:
                fp.write(image_data)
            n = n + 1


if __name__ == '__main__':
    keyword = '垃圾桶'
    save_dir = './all'
    page_num = 2
    get_images_from_baidu(keyword, page_num, save_dir)
    print('Get images finished.')

1.3 代码使用说明

1. 在keyword 处设置想要检索的图片，以文字为描述
2. save_dir 是爬取下来的图片要保存的具体位置
3. page_num 控制的是爬取的图片数量，图片数量是page_num*30
4. 爬取下来的图片会从0开始排序，以JPG格式保存

2. 数据标注

2.1 标注工具的安装

首先假设已经安装了Anaconda，如果没安装的，可以看之前写的一篇关于深度学习环境配置的一篇博客，链接在这windows下安装anaconda3（对应的python版本为3.8）+pytorch1.6.0。

使用的是labelimg的图像数据标注工具，安装步骤如下：

1. 创建虚拟环境
conda create -n labelimg python=3.7
2. 激活虚拟环境
conda activate labelimg
3. 安装依赖
pip install PyQt5
pip install pyqt5-tools
pip install lxml
4. 安装labelimg
pip install labelimg

2.2 labelimg使用

1. 在labelimg环境下，命令行输入labelimg即可打开，界面如下：
   
2. labelimg常用快捷键
   
3. 打开爬取图片的保存文件夹
   打开文件夹
   

首先显示第一张图片，右下角是文件列表，以文件名排序

4. 选择YOLO模式
共有三种模式供选择，显示的是什么就是什么模式

YOLO模式

其它模式

5. 标注
   快捷键W创建矩形框并给出是否满溢的标签，图中为满溢，因此给出overflow的标签
   
   点击“OK”会在右侧给出标签信息
   
   已经给出过的标签会被记录，不需要重复输入，选择即可
   
6. 在标注完一张图之后，CTRL + S保存后，进行下一张的标注即可，每张图像单独保存一个标签文件，具体见图
   
   第一列是代表的类别，其余四列是矩形框的位置
   由于展示的示例图中仅存在满溢的情况，所以类别为0，即overflow

3. YOLO代码准备

3.1 代码下载

1. 下载代码
git clone https://github.com/ultralytics/yolov5  # clone
2. 环境安装
cd yolov5
pip install -r requirements.txt  # install
3. 说明：在使用时我重新建立了一个名为yolo的虚拟环境来安装yolo所需的环境

3.2 创建自己的参数配置文件

以COCO128.yaml文件为基础进行了修改，具体如下

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/garbage  # dataset root dir
# 训练和测试在这里使用了相同的数据
train: C:\Users\39205\Desktop\smartcity\yolov5\datasets\garbage\images\train  # train images (relative to 'path') 128 images
val: C:\Users\39205\Desktop\smartcity\yolov5\datasets\garbage\images\train  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: overflow
  1: Not full


# Download script/URL (optional)
#download: https://ultralytics.com/assets/coco128.zip

3.3 训练

使用train.py脚本进行训练，参数传入如下：

--img
640
--epochs
3
--data
garbage.yaml
--weights
yolov5s.pt

3.4 实验数据保存

1.实验数据保存在根目录下，具体如下：

runs\train\exp

2.部分检测结果示例如下

3.训练曲线如下

4.标签情况如下

4. 对单张图像或者视频进行检测

4.1 使用detect.py，参数如下

# 自己训练的权重
	parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'weights/best.pt', help='model path or triton URL')
# 测试的图像路径
    parser.add_argument('--source', type=str, default=ROOT / 'testdata/000000.jpg', help='file/dir/URL/glob/screen/0(webcam)')
# 之后的基本与源代码一致
    parser.add_argument('--data', type=str, default=ROOT / 'data/garbage.yaml', help='(optional) dataset.yaml path')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='show results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--visualize', action='store_true', help='visualize features')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')

4.2 测试图像输入

命令行传入参数

python detect.py --source ./testdata/img/000000.jpg

测试图像如下：

测试结果如下：

4.3 测试视频输入

命令行传入参数

python detect.py --source ./testdata/video/vid.mp4

视频测试过程：

测试视频中的部分帧

测试结果：

5. 总结

可以检测到垃圾箱，但是可视化的结果其实并不是很好，这可能是因为仅使用了400多张数据进行训练，还有一部分原因可能是epoch数量设置的比较少，仅仅训练了3个epoch，不过基本可以检测到垃圾箱，后续就是拿到服务器上训练，并进一步增加数据量了。并且在标注过程中可能并不是特别精细，下一步实施的时候也可以注意改进这一部分。总体来说，这一套流程是走通了，不过大规模的数据标注还是很麻烦，最近在研究自动标注的东西，欢迎交流讨论。