笨比(我)训练yolo
在上一次配置好ubuntu端的rknn环境后,我们现在需要训练一个检测器,由于项目需要,我们选择YOLO
数据集就更简单了,常见的数据集都可以,为了方便后续,选择:红外航拍人车检测数据集,光电红外开源平台
下载数据集
数据集结构如下:
datasets
---dataset
----00000.jpg
----00001.jpg
-----…………
-----00001.xml
-----00002.xml
………………
---readme.txt
我们需要将数据集分开 ,并移动到代码的data_dir中
classes
根据数据集中的readme,数据集分为6类,
classes = ["person","car" ,"bus","cyclist","bike","truck"]
Stupid Methods
为了让小白也可以开train,我们对数据集的处理尽量使用笨方法:
第一步 将所有的jpg的路径打包为txt
create_txt.py
# -*- coding:utf-8 -*-
import glob
imageList = glob.glob(r"reddata\data\dataset\*.jpg") # 图片所在文件夹的路径
f = open('train.txt', 'a') # 创建标签文件存放图片文件名
for item in imageList:
print(item)
img_name = item.split('/')[-1] # 图片文件名018.jpg
# img_name = item.split('\\')[-1] # 图片文件名018.jpg
f.write(img_name + '\n') # 将图片文件名逐行写入txt
f.close()
print('OK')
#所有的jpg文件
将所有的xml文件转为labels
xml2labels.py
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = []
classes = ["person","car" ,"bus","cyclist","bike","truck"] ##修改为自己的类别
#原样保留。size为图片大小
# 将ROI的坐标转换为yolo需要的坐标
# size是图片的w和h
# box里保存的是ROI的坐标(x,y的最大值和最小值)
# 返回值为ROI中心点相对于图片大小的比例坐标,和ROI的w、h相对于图片大小的比例
def convert(size, box):
dw = 1./(size[0])
dh = 1./(size[1])
x = (box[0] + box[1])/2.0 - 1
y = (box[2] + box[3])/2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def convert_annotation(image_add):
#image_add进来的是带地址的.jpg
image_add = os.path.split(image_add)[1] #截取文件名带后缀
image_add = image_add[0:image_add.find('.',1)] #删除后缀,现在只有文件名没有后缀
#现在传进来的只有图片名没有后缀
in_file = open('reddata/data/xml/' + image_add + '.xml',encoding='utf8', errors="ignore") #修改为你自己的输入目录
out_file = open('reddata/data/labels/%s.txt'%(image_add), 'w',encoding='utf8', errors="ignore") #修改为你自己的输出目录
tree=ET.parse(in_file)
root = tree.getroot()
if root.find('size'):
size = root.find('size')
w = int(size.find('width').text) #偶尔xml标记出错,width或height设置为0了
h = int(size.find('height').text) #需要标记出来,便于单独处理
if w==0:
print("出错! width或height为0: "+image_add)
return
#在一个XML中每个Object的迭代
for obj in root.iter('object'):
#iter()方法可以递归遍历元素/树的所有子元素
difficult = obj.find('difficult').text
cls = obj.find('name').text
#如果训练标签中的品种不在程序预定品种,或者difficult = 1,跳过此object
if cls not in classes or int(difficult)==1:
continue
#cls_id 只等于1
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
#b是每个Object中,一个bndbox上下左右像素的元组
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
bb = convert((w,h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
else:
print("出错!xml缺少size: "+image_add) #偶尔xml缺少size,需要标记出来,便于单独处理
image_adds = open("train.txt",encoding='utf8') #修改为你自己的训练数据集目录
for image_add in image_adds:
print(image_add)
image_add = image_add.strip()
print (image_add)
convert_annotation(image_add)
翻看yolo代码,发现他的数据集格式是这样的
datasets
---test
----images
----00056.jpg
----labels
----00056.txt
……
---train
----images
----00073.jpg
----labels
----00073.txt
……
----val
----images
----02752.jpg
----labels
----02752.txt
……
我们仿照他的结构
根据训练集,测试集,验证集的比例将数据集划分
data_split.py
import os, random, shutil
def moveFile(fileDir):
pathDir = os.listdir(fileDir) #取图片的原始路径
filenumber=len(pathDir)
rate=1 #自定义抽取图片的比例,比方说100张抽10张,那就是0.1
picknumber=int(filenumber*rate) #按照rate比例从文件夹中取一定数量图片
sample = random.sample(pathDir, picknumber) #随机选取picknumber数量的样本图片
print (sample)
for name in sample:
shutil.move(fileDir+name, tarDir+name)
return
if __name__ == '__main__':
fileDir = "./newdatastes/reddata/data/dataset/" #源图片文件夹路径
tarDir = './datasets/val/images/' #移动到新的文件夹路径
moveFile(fileDir)
这里可以根据自己想要的比例抽取数据集,需要注意的是:运行代码后已经移动的jpg无法恢复
根据已经存在的train,test,val中的 jpg 将相应的labels 移动过来
由于我们是随机抽取的jpg数据集,所以我们不能直接运行上面的代码继续移动labels
所以我们先让数据集产生txt目录,然后根据这些索引去移动对应的labels:
name2txt.py
import os
file_path = "./val/images/" # Path
path_list = os.listdir(file_path)
path_name = []
for i in path_list:
path_name.append(i.split(".")[0])
#path_name.sort()
for file_name in path_name:
with open("val.txt", "a") as file:
file.write(file_name +".txt" + "\n")
print(file_name)
file.close()
分别将三个文件夹下的jpg目录产出:
train.txt
test.txt
val.txt
根据txt文件移动对应的labels
movetxt.py
import shutil
if __name__ == '__main__':
file_object = open('./datasets/val.txt')
try:
for line in file_object:
# print(line)
shutil.move('./newdatastes/reddata/data/labels/'+line.rstrip('\n'), "./datasets/val/labels")
finally:
file_object.close()
注意不要有中文路径
成功制作数据集!!
其实用另一种方法可以很快速的制作出这种数据集,但是在前两天有个学弟问过我一些制作数据集的问题后,我意识到,一些没做过这方面的人可能更倾向于笨方法去处理,我也很笨hh,所以还是写出来了
更改yolo 代码
更改data/xxxx.yaml
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ./datasets # dataset root dir
train: ./datasets/train/images # train images (relative to 'path') 128 images
val: ./datasets/val/images #images/train2017 # val images (relative to 'path') 128 images
test: ./datasets/test/images # test images (optional)
# number of classes
nc: 6 #更改为你自己的classesnum
# Classes
names: [ 'person', 'car', 'bus', 'cyclist', 'bike', 'truck']#更改为你自己的classes name
其他的就没什么大的修改点了,有什么bug就去查好了(doge)
记得install requirements
我是用的yolo 代码依赖于1.8.0的torch,但是torch版本
但是有GPU 不用是笨蛋hh
所以我安装了1.8.1touch+cu111
记得去看看你们的cuda是否可用!!!
遇事不决 train一发
先conda 激活环境
conda activate env
然后既可以开始train,根据你自己的设备情况选择epochs和batch-size
python train.py --epochs 300 --batch-sie 32 --device 0
经过十个小时的训练,已经成功训练完毕,爱训练过程中可以利用tensorboard 去看看 训练的收敛情况,但是我懒hh
如果不出意外的话,训练好的pt文件会保存在runs/train/exp_x/weights/best.pt
在weights 中还保留有训练中产生的各种参数:
such as :
还有其他的就不一一展示了
对于笨比(我)来说,看曲线不如detect
所以:
python detect.py --weights ./runs/train/weights/best.pt --source ./reference/images -devie 0
可以看到速度还是很快的:
效果也很不错