背景
本项目背景是使用一个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 shutildef convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[2] / 2.0) * dwy = (box[1] + box[3] / 2.0) * dhw = box[2] * dwh = box[3] * dhreturn (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的类别IDsrc_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 = 0for ann in data['annotations']:if ann['image_id'] == img_id and ann['category_id'] == category_id:bHasId = 1breakif 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 + filenamedst_path = dst_img_path + filenameshutil.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 labelssegments = False # segment or box labelsdir = 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
