基于yolov5只进行人物目标检测

背景

本项目背景是使用一个usb摄像头模组，在一定距离下检测人，并判断人是否进入设定的区域中。因此，需要一种目标检测模型，我们选择yolov5作为检测网络，把原来包含80类的coco数据集提取出包含人物的图像，重新训练一次。
我使用的yolov5的代码版本如下图框选所示

人物数据集提取

在coco网站下载coco2017训练、验证、测试数据集及其标签数据
coco数据下载

使用如下python的代码提取人物数据及其标签到新的文件夹中，因为coco2017数据集中的标签格式为.json，需要将其中的人物标签信息获取，并保存为.txt格式，txt中每一行的格式为
[标签ID] [标记框左上角x坐标] [标记框左上角y 坐标] [标记框宽度] [标记框高度]
共5个元素
值得注意的是，如果想要尽量避免修改读取模型的代码，就和train.py中提到的coco128数据文件格式相同，当前是
coco128
|----images/
|— images/train2017/
|----labels/
|----labels/train2017/
|-----LICENSE
|-----README.txt
因此，我的文件格式为
coco2017_person
|----images/
|— images/train2017/
|— images/val2017/
|----labels/
|----labels/train2017/
|----labels/val2017/

import os
import json
import shutil

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[2] / 2.0) * dw
    y = (box[1] + box[3] / 2.0) * dh
    w = box[2] * dw
    h = box[3] * dh
    return (x, y, w, h)

def main():
    json_file = 'E:/pycharm/coco2017/annotations_trainval2017/annotations/instances_train2017.json'  # 替换为你的JSON文件路径
    save_path = 'E:/pycharm/coco2017_person/labels/train2017/'  # 替换为你要保存YOLO格式标签的路径
    category_id = 1  # COCO中person的类别ID
    src_img_path = 'E:/pycharm/coco2017/train2017/' #替换为你的原始图像路径
    dst_img_path = 'E:/pycharm/coco2017_person/images/train2017/' #替换为你的保存的图像路径

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for img in data['images']:
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"
        bHasId = 0
        for ann in data['annotations']:
            if ann['image_id'] == img_id and ann['category_id'] == category_id:
                bHasId = 1
                break

        if bHasId:
            f_txt = open(os.path.join(save_path, ana_txt_name), 'w')
            for ann in data['annotations']:
                if ann['image_id'] == img_id and ann['category_id'] == category_id:
                    box = convert((img_width, img_height), ann["bbox"])
                    f_txt.write("0 %s %s %s %s\n" % (box[0], box[1], box[2], box[3]))
            f_txt.close()
            src_path = src_img_path + filename
            dst_path = dst_img_path + filename
            shutil.copy(src_path, dst_path)

if __name__ == '__main__':
    main()

提取后的数据都与人物相关

使用仅含人物的数据集训练

在yolov5\data\文件夹下面，复制一份coco128.yaml, 并修改名字为coco_person.yaml

# coco_person.yaml
# Ultralytics YOLOv5 🚀, AGPL-3.0 license
# COCO 2017 dataset http://cocodataset.org by Microsoft
# Example usage: python train.py --data coco.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco  ← downloads here (20.1 GB)

# 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: ../coco2017_person # dataset root dir
train: images/train2017 # train images (relative to 'path') 118287 images
val: images/val2017 # val images (relative to 'path') 5000 images
test: # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

# Classes
names:
  0: person

# Download script/URL (optional)
download: |
  from utils.general import download, Path


  # Download labels
  segments = False  # segment or box labels
  dir = Path(yaml['path'])  # dataset root dir
 # url = 'https://github.com/ultralytics/assets/releases/download/v0.0.0/'
 # urls = [url + ('coco2017labels-segments.zip' if segments else 'coco2017labels.zip')]  # labels
 # download(urls, dir=dir.parent)

  # Download data
 # urls = ['http://images.cocodataset.org/zips/train2017.zip',  # 19G, 118k images
 #         'http://images.cocodataset.org/zips/val2017.zip',  # 1G, 5k images
 #         'http://images.cocodataset.org/zips/test2017.zip']  # 7G, 41k images (optional)
 # download(urls, dir=dir / 'images', threads=3)

执行命令训练

python train.py --data coco_person.yaml --weights yolov5s.pt --img 640

测试结果
原始数据的训练结果验证yolov5s.pt

python detect.py --weights yolov5s.pt --source E:\pycharm\yolov5\data\images

仅含人物数据的训练结果验证E:\pycharm\coco2017_person\exp9\weights\best.pt，可以看到风筝、书包等不再检测，仅检测人

python detect.py --weights E:\pycharm\coco2017_person\exp9\weights\best.pt --source E:\pycharm\yolov5\data\images