yolov5 设置只检测某几个固定的类

下载 yolov5 代码

yolov5 代码下载（我选的是 master 分支，其他如 v6.1 应该也可以）：https://github.com/ultralytics/yolov5

问题描述

如果使用 coco 数据集进行检测，coco 数据集中共有 80 个类别，但我只想关注前三个类别 ‘person’, ‘bicycle’, ‘car’，其他类别不关注。
找到 data/coco.yaml， ‘person’, ‘bicycle’, ‘car’ 这三个类对应的下标为 0, 1, 2

# Classes
nc: 80  # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']  # class names

问题解决

修改代码

打开 utils/datasets.py，搜索include_class，找到下面这段代码的位置，大约在 450 行的位置：

修改代码为：

include_class = [0, 1, 2]	# 因为我只想关注 'person', 'bicycle', 'car'，这三个类别在 names 中的位置是 0, 1, 2

亲测 v6.1 版本也可以找到 include_class

然后运行 val.py 或 train.py 即可。