yolov5在kaggle的水稻检测表现很是好,至少是单模型我所知道的最高得分。对比mmde双阶段模型和efficientdet模型,上一篇文章讲解了pytorch最强复现的yolo4版本,yolov5基本和yolo4训练过程同样。 https://github.com/ultralytics/yolov5python
第一步先生成yolo统一的格式txt文件,
git
import os name=os.listdir('./image') for i in range(len(name)): name[i]='./coco/images/train2017/'+name[i] file = open('./train2017.txt', 'w') for i in range(len(name)): file.write(name[i]) file.write('\n') file.close()
生成train2017是遍历你本身的图像名称,一共有二个文件,分为训练集和验证集,我直接全训练集训练了,随便复制一点路径在生成一个testdev.txt文件便可。github
生成一堆txt文件,每一个对应一个xml文件生成。
yolo统一格式label xywh形式。web
import xml.etree.ElementTree as ET import pickle import os from os import listdir, getcwd from os.path import join sets = [] classes = ["wheat"] # 原样保留。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)] # 删除后缀,如今只有文件名没有后缀 print(image_add) # 如今传进来的只有图片名没有后缀 in_file = open('./xml/' + image_add + '.xml') out_file = open('./label/%s.txt' % (image_add), 'w') tree = ET.parse(in_file) root = tree.getroot() size = root.find('size') w = int(size.find('width').text) h = int(size.find('height').text) # 在一个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') image_adds=os.listdir('./image') for image_add in image_adds: # print(image_add) image_add = image_add.strip() # print (image_add) convert_annotation(image_add)
主须要二处须要修改,image是个人图像目录,xml文件夹下是我本身的xml文件存放处,label是生成txt文件的位置。svg
修改data下的coco.yaml,类别改为本身数据集上的数量,和生成的txt文件。
修改model下的yaml文件,将类别数量修改为本身数据集的类别
布局
个人数据是从yolov4直接拉过来的因此有个地方是没有用的,和images,lables同级的txt文件是没有用的。训练yolov5调用的是coco/train2017.txt文件。
我这里所有修改为了当前路径下的./coco,恰好对应我xml转txt的文件。
train2017.txt也是设置的当前路径下的coco训练集,这样路径就不会出错,本身检查一下就行ui
!python train.py --data coco.yaml --cfg yolov5s.yaml --weights './weights/yolov5s.pt' --batch-size 5
这里的pt预训练权重是在google colab下载的,google
cpu下的推理速度也挺快,
spa