yolov5 检测一类物体

文章目录[隐藏]

使用yolov5官方框架检测一类物体

使用yolov5官方框架检测一类物体

yolov5的官方框架可较好的对共80种类进行目标检测，本文介绍一种直接修改源代码来只检测一类物体的方法以及通用的方法（利用数据集训练自己的权重）。

一、直接修改源代码
1.在GitHub上下载官方代码。https://github.com/ultralytics/yolov5。

2.对环境进行配置，安装requirements.txt进行配置，可以在pychram的terminal键入
pip install -r requirements.txt
会自动进行环境配置，使用yolov5大多数出现的问题都是因为环境配置不正确。

3.利用官方权重测试。

parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam**

主要对detect.py中的weights和source进行修改，这里weights选择四个官方权重中的yolov5s.pt，source即测试数据的来源。当无附加命令时，使用default中的参数。

4.测试结果：
在这里插入图片描述
输出结果在runs的detect中。可以看到官方权重可以检测人、领带等物体。

# number of classes
nc: 80

# class names
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' ]

yolov5默认权重共可检测以上80种类。

五、代码修改。（以只检测人为例。）
将detect.py中98行开始修改为如下代码，加入了一个判断语句。即检测物体为人时才打印结果，其他类不打印。

# Print results
                for c in det[:, -1].unique():
                    if names[int(c)] == 'person':
                      n = (det[:, -1] == c).sum()  # detections per class
                      s += f'{n} {names[int(c)]}s, '  # add to string
                    else:
                        continue

106行开始修改为如下代码。

# Write results
    for *xyxy, conf, cls in reversed(det):
        if save_img or view_img:  # Add bbox to image
            if names[int(cls)] == 'person':
               label = f'{names[int(cls)]} {conf:.2f}'
               plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
                    else:
                        continue