YOLOv5 是COCO数据集上预处理的一系列对象检测架构和模型,代表Ultralytics对未来视觉人工智能方法的开源研究,融合了数千小时研究和开发过程中积累的经验教训和最佳实践。
本文用来记录第一次使用 YOLOv5实现:
博主所使用的环境是win10 + Anaconda的虚拟环境(Python>=3.7.0,PyTorch>=1.7),软件环境是Jupyter Notebook
提示:以下是本篇文章正文内容,下面案例可供参考
下载地址:https://github.com/KZwyj/yolov5
使用YOLOv5的模型可以识别以下物体类别等:
'dog','person','tie','cat','tv','car','meatballs','marinara sauce','tomato soup','chicken noodle soup','french onion soup','chicken breast','ribs','pulled pork',''cavity'
同时也可以自定义数据集并训练,有关于自定义数据集这部分后期再进行讲解。
基本环境要求:
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.1
Pillow>=7.1.2
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.64.0
其他(不安装也不会影响运行):
tensorboard>=2.4.1
pandas>=1.1.4
seaborn>=0.11.0
ipython # interactive notebook
psutil # system utilization
thop>=0.1.1 # FLOPs computation
在虚拟环境中安装以上库:
cd yolov5
pip install -r requirements.txt # install
测试YOLOv5:
model = torch.hub.load("ultralytics/yolov5","yolov5s")
出现Adding AutoShape…提示,即为YOLOv5模型(yolov5s)加载成功:
然后进行测试(图片目标检测):
img = "https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/zidane.jpg"
results = model(img)
results.print()
显示图片:
%matplotlib inline
# squeeze从数组的形状中删除单维度条目,即把shape中为1的维度去掉
plt.imshow(np.squeeze(results.render()))
plt.show()
从前面我们知道,我们已经获取到了YOLOv5模型,下面使用opencv,加载视频并进行目标检测。
import torch
from matplotlib import pyplot as plt
import numpy as np
import cv2
video_path = "./videoes/traffic_evening.mp4"
cap = cv2.VideoCapture(video_path)
我所加载的视频是放在当前目录的videoes文件夹下,可以根据需要修改video_path为你视频存放位置。
while cap.isOpened():
# 参数ret 为True 或者False,代表有没有读取到图片帧
# 第二个参数frame表示截取到一帧的图片
ret,frame = cap.read()
# 将图片帧输入模型
results = model(frame)
# 输出目标检测结果到屏幕
cv2.imshow('YOLO',np.squeeze(results.render()))
# 退出
if cv2.waitKey(10) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
摄像头目标检测和视频目标检测区别不大,只需要修改:
cap = cv2.VideoCapture(0)