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

下载 yolov5 代码

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

问题描述

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

# 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，找到下面这段代码的位置，大概在 470 行左右：

修改代码为：

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

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