安装
-
先把项目下载到本地
git clone https://github.com/ultralytics/yolov5.git
也可以到网站下载压缩包解压
-
我用的原生python,希望创造一个独立的虚拟环境,所以这里用到了
virtualenv
,在你的工作运行python -m venv yolo
会在当前文件夹创建一个纯净的python环境,以后下载的所有包都在yolo文件夹里,而不会和你原生环境的包版本冲突,当然你也可以跳过这一步,不创建虚拟环境,直接在原生环境安装依赖。
运行成功后进入yolo文件夹下的Scripts文件夹,在命令行运行
activate.bat
激活虚拟环境,如图此时你的python、pip都是执行在虚拟环境里了。
-
将下载好的yolov5源码文件夹移入Scripts文件夹,进入yolov5文件夹,运行
pip install -r requirements.txt
由于默认安装的pytorch是cpu版本的,为了提高运行速度,这里我们重新安装一个gpu版本的pytorch
pip install torch==1.9.1+cu102 torchvision==0.10.1+cu102 torchaudio===0.9.1 -f https://download.pytorch.org/whl/torch_stable.html
在安装pytorch-gpu的时候可能会遇到没有cuda的情况,安装cuda可以自己去网上查找教学,这里不说明了,记得更改对应上面的cuda版本号,不过我本机安装的cuda是10.0,运行上面的cuda10.2也运行成功了。
运行试试
1、进入到yolov5,运行
python detect.py
第一次运行会自动从网上下载最新的yolov5s.pt
权重
运行结束后会在当前文件夹生成runs\detect\exp
目录,里面会利用模型识别出的两张示例图片
同时根据官网示例,还可以用相同的命令识别不同的媒体
python detect.py --source 0 # webcam
file.jpg # image
file.mp4 # video
path/ # directory
path/*.jpg # glob
'https://youtu.be/NUsoVlDFqZg' # YouTube
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
为了更清晰为啥会得这样的结果,下面进入detect.py
文件看一些关键参数
–weights 是使用到的权重文件
–source 是使用的源文件
–img-size 是图片大小,也就是讲源文件压缩到640*640的大小
–device 是使用到的GPU设备,当你有多个GPU时可以自定义编号,这里pytorch会自动探测取0编号GPU
–project 识别结果的输出文件夹
训练自己的模型
Github上面有yolov5官方的训练教学,用到的是COCO数据集,这里是训练自己的数据集
创建数据集
首先需要知道的是,如果要得到一个稍微好一点的模型,需要至少200张图片,当然你测试的话可以少一点,
1、下载打标签软件https://github.com/tzutalin/labelImg,用法自己了解了,这一步是体力活了,将所有图片打上你要识别的标签,创建一个datasets文件夹,在里面创建annotations
、images
、labels
文件夹,结构如图
images保存用于打标签的图片,annotations保存打好标签的xml文件,labels文件夹先不用管
2、由于yolo训练用到的是txt格式,所以需要转换xml到txt,编写transferm.py
脚本文件,代码如下
import os
import xml.etree.ElementTree as ET
def ConverCoordinate(imgshape, bbox):
# 将xml像素坐标转换为txt归一化后的坐标
xmin, xmax, ymin, ymax = bbox
width = imgshape[0]
height = imgshape[1]
dw = 1. / width
dh = 1. / height
x = (xmin + xmax) / 2.0
y = (ymin + ymax) / 2.0
w = xmax - xmin
h = ymax - ymin
# 归一化
x = x * dw
y = y * dh
w = w * dw
h = h * dh
return (x,y,w,h)
def readxml(filename):
outfile = open('labels/{}.txt'.format(filename), 'w')
filetree = ET.parse('annotations/{}.xml'.format(filename))
root = filetree.getroot()
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
imgshape = (width, height)
for obj in root.findall('object'):
# 获取类别名,判断是否在classes中,不存在则跳过。
obj_name = obj.find('name').text
if obj_name not in classes:
continue
obj_id = classes.index(obj_name)
# 获取每个obj的bbox框的左上和右下坐标
bbox = obj.find('bndbox')
xmin = float(bbox.find('xmin').text)
xmax = float(bbox.find('xmax').text)
ymin = float(bbox.find('ymin').text)
ymax = float(bbox.find('ymax').text)
bbox_coor = (xmin, xmax, ymin, ymax)
txtvalue = ConverCoordinate(imgshape, bbox_coor)
outfile.write('{}'.format(obj_id) + ' ' + ' '.join([str(i) for i in txtvalue]) + '\n')
if __name__ == '__main__':
# 超参数
classes = ['hero','flash','skill','restore']
# 配置JPEG文件路径
localdir = os.getcwd()
Imagesfiledir = os.path.join(localdir, 'images')
for filename in os.listdir(Imagesfiledir):
readxml(filename[:-4])
注意!!!
上面classes保存的是你打的所有标签,顺序很重要。
3、此时还要编写一个文件,以至于符合yolov5训练需要的数据集格式,也就是划分数据集,编写split.py
,键入如下代码:
import os
import shutil
import random,math
src = 'images'
dst = 'labels'
img_name = os.listdir('./images')
amount = len(img_name)
img_format = img_name[0][-4:]
img_name = [i[:-4] for i in img_name]
length = math.floor(0.8 * amount)
print(length)
train_labels_path = 'train/labels'
train_images_path = 'train/images'
val_images_path = 'val/images'
val_labels_path = 'val/labels'
# 创建文件夹
if not os.path.exists(train_labels_path):
os.makedirs(train_labels_path)
if not os.path.exists(train_images_path):
os.makedirs(train_images_path)
if not os.path.exists(val_images_path):
os.makedirs(val_images_path)
if not os.path.exists(val_labels_path):
os.makedirs(val_labels_path)
# 创建训练集
for i in range(length):
img = random.choice(img_name)
images = os.path.join(src,img+img_format)
print(images)
labels = os.path.join(dst,img+'.txt')
shutil.move(images,train_images_path)
shutil.move(labels,train_labels_path)
img_name.remove(img)
# 创建测试集
for i in range(amount-length):
images = os.path.join(src,img_name[i]+img_format)
labels = os.path.join(dst,img_name[i]+'.txt')
shutil.move(images,val_images_path)
shutil.move(labels,val_labels_path)
将两个脚本文件放在和三个文件夹同级的目录下,依次运行后,得到如下目录结构
此时可以删掉除train、val的其他文件和文件夹了,即剩下
终于得到数据集了。。。
修改配置
1、创建自己的配置文件hero.yaml
train: ./data/datasets/train/images
val: ./data/datasets/val/images
nc: 4 # 下面标签的数量
names: ['hero','flash','skill','restore']
将hero.yaml和datasets文件夹都放入yolov5目录下的data文件夹
2、修改train.py文件,如图,修改几个关键参数
–weights 初始权重,这个默认yolov5s,可以不加,但加上效果可能好一点
–data 配置文件名,需要更改为自己的datasets
–epochs 训练世代,默认是300
–batch-size 批训练大小,需要加载进内存,16可能有点大,这里根据自己内存大小动态调整
–img-size 统一压缩到的图片大小,这里默认就好
实际上只需要简单改一下–data
3、训练试试
python train.py
等待训练完成即可
如果出现错误如下
那么可以改小–batch-size的大小
或者增加内存大小,方法:https://blog.csdn.net/weixin_43817670/article/details/116748349