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
至此完成只检测一类物品的任务。
总结:这是一种讨巧的方法,实际上模型仍然对80种类进行了检测,只是没有显示和打印出结果。相较于自己训练的单个类的模型,这种方法在实时性上比较差。