写文章Yolov5实践龟壳人工智能方向研究生27 人赞同了该文章

YOLOv5实践

20201009

最近在做目标检测相关的任务，最终是要做落地的产品，部署到移动端NPU上，刚开始尝试了yolov4,发现训练代价很大，最终模型240M左右，有点大，于是开始尝试yolov5，发现yolov5太香了，训练快，精度高，模型小。

参考网址：

github.com

YOLOv5 COCO数据集 训练 |简记​

blog.csdn.net

环境：Ubuntu16.04 + Pytorch1.6.0 + Python3.8.5 + CUDA10.1

一、代码下载以及环境配置

将代码clone下来按照教程安装配置好环境以及需要的库就可以了，教程种给的是coco128的数据集，可以自己先跑一下，保证自己的整体环境及代码没有问题。

二、训练自己的数据集

2.1 数据处理

我的原始数据集是公司自己采集和打标签的，格式是文件夹中包含图片和对应的xml文件，如图1所示：

由于图片的格式有jpg,JPG,还有jpeg,图片的名字标注也是各式各样，我就将其按照数字1-n统一修改了一下，图片格式统一保存为.jpg格式，代码如下：

modify_image_and_xml.py

 code by zzg 2020-05-30#!/usr/bin/python# -*- coding: UTF-8 -*-# get annotation object bndbox location import xml.etree.cElementTree as ET except ImportError: import xml.etree.ElementTree as ETimport os,sysimport globimport cv2from PIL import Imageimport matplotlib.pyplot as pltimport pdb#the direction/path of Image,Labelsrc_img_dir = image_src src_xml_dir = xml_src dst_img_dir = image dst_xml_dir = xml if not os.path.exists(dst_img_dir): os.makedirs(dst_img_dir)if not os.path.exists(dst_xml_dir): os.makedirs(dst_xml_dir)img_Lists = glob.glob(src_img_dir + /*.JPG )#print(img_Lists)# rename = src_img_dir.split( / )[-1]# print(rename)img_basenames = []for item in img_Lists: img_basenames.append(os.path.basename(item)) #print(img_basenames)img_name = []for item in img_basenames: temp1, temp2 = os.path.splitext(item) img_name.append(temp1)# print(img_name)cnt = 0for img in img_name: im = cv2.imread(src_img_dir + / + img + .JPG ) # print(type(im)) print(im.shape[::-1]) channels, width, height = im.shape[::-1] ##get w and h ##read the scr_xml AnotPath = src_xml_dir + / + img + .xml  tree = ET.ElementTree(file=AnotPath)  root = tree.getroot() ObjectSet = root.findall( object ) ObjBndBoxSet = [] ObjBndBoxSet1 = {}  for Object in ObjectSet: ObjName = Object.find( name ).text BndBox = Object.find( bndbox ) x1 = int(BndBox.find( xmin ).text)#-1  y1 = int(BndBox.find( ymin ).text)#-1 x2 = int(BndBox.find( xmax ).text)#-1 y2 = int(BndBox.find( ymax ).text)#-1 BndBoxLoc = [ObjName,x1,y1,x2,y2] # print(x1,y1,x2,y2) ObjBndBoxSet.append(BndBoxLoc)  print(ObjBndBoxSet) # save the crop-image in dst_crop cnt += 1 cv2.imwrite(dst_img_dir + / + str(cnt) + .jpg , im)  # rewrite xml to dst_xml xml = open((dst_xml_dir + / + str(cnt) + .xml ), w ) xml.write( annotation \\n ) xml.write( \\t folder + VOC2007 + /folder \\n ) xml.write( \\t filename + str(cnt)+ .jpg + /filename \\n ) xml.write( \\t source \\n ) xml.write( \\t\\t database Unknown /database \\n ) xml.write( \\t /source \\n ) xml.write( \\t size \\n ) xml.write( \\t\\t width + str(width) + /width \\n ) xml.write( \\t\\t height + str(height) + /height \\n ) xml.write( \\t\\t depth + str(channels) + /depth \\n ) xml.write( \\t /size \\n ) xml.write( \\t\\t segmented 0 /segmented \\n ) print( ===========start rewrite bndbox============== ) for x in ObjBndBoxSet: # print(x) [classname,x1,y1,x2,y2] = x  xml.write( \\t object \\n ) xml.write( \\t\\t name + classname + /name \\n ) xml.write( \\t\\t pose Unspecified /pose \\n ) xml.write( \\t\\t truncated 1 /truncated \\n ) xml.write( \\t\\t difficult 0 /difficult \\n ) xml.write( \\t\\t bndbox \\n ) xml.write( \\t\\t\\t xmin + str(x1) + /xmin \\n ) xml.write( \\t\\t\\t ymin + str(y1) + /ymin \\n ) xml.write( \\t\\t\\t xmax + str(x2) + /xmax \\n ) xml.write( \\t\\t\\t ymax + str(y2) + /ymax \\n ) xml.write( \\t\\t /bndbox \\n ) xml.write( \\t /object \\n )  xml.write( /annotation ) print(cnt)print( =======================finished!=================== )

之后又将xml标签和图片分成xml和image这2个文件夹，并且制作训练集和测试集，生成train.txt和test.txt，代码如下：

01. split_xml_jpg.py

 code by zzg -2020-10-07#复制或者移动一个文件夹下的所有图片或者其他指定文件到另一个文件夹import osimport shutilpath = data/smoke/ new_path = image/ new_path1 = xml/ count = 0for root, dirs, files in os.walk(path): for i in range(len(files)): #if (files[i][-3:] == jpg or files[i][-3:] == JPG ): if (files[i][-3:] == jpg ): count += 1 file_path = root + / + files[i] new_file_path = new_path + / + files[i] shutil.copy(file_path, new_file_path) #shutil.move(file_path, new_file_path))print(count)print( move finished!! )

02. make_train_val_set.py

 code by zzg-2020-10-07# -*- coding: utf-8 -*-import osimport random#first create# os.makedirs( VOC2007/Annotations )# os.makedirs( VOC2007/ImageSets )# os.makedirs( VOC2007/ImageSets/Main )# os.makedirs( VOC2007/JPEGImages )def _main(): val_percent = 0.2  train_percent = 0.8 xmlfilepath = xml  total_xml = os.listdir(xmlfilepath) num = len(total_xml) list = range(num) tv = int(num * val_percent) tr = int(num * train_percent) val = random.sample(list, tv) ftest = open( VOC2007/ImageSets/Main/test.txt , w ) ftrain = open( VOC2007/ImageSets/Main/train.txt , w ) for i in list: name = total_xml[i][:-4] + \\n  if i in val: ftest.write(name) else: ftrain.write(name) ftest.close() ftrain.close()if __name__ == __main__ : _main()

之后制作yolov5所需的数据格式，需要自己写一个classes.names的文件，将自己的类别放进去，我这里是通过xml转json,之后把json和image放到同一个文件夹，然后调用脚本来转换，还可以自己聚类自己数据集的anchor：

01.xml_json.py

 code by zzg 2020-05-13##xml_to_jsonimport globimport xmltodictimport jsonpath = /workspace/zigangzhao/yolov5/tool/dms/xml_modify/ path2 = /workspace/zigangzhao/yolov5/tool/dms/json/ xml_dir = glob.glob(path + *.xml )print(xml_dir)def pythonXmlToJson(path): xml_dir = glob.glob(path + *.xml ) # print(len(xml_dir)) for x in xml_dir: with open(x) as fd: convertedDict = xmltodict.parse(fd.read()) jsonStr = json.dumps(convertedDict, indent=1) print( jsonStr= ,jsonStr) print(x.split( . )[0]) json_file = x.split( . )[0].split( / )[-1] + .json  with open(path2 + / + json_file, w ) as json_file: json_file.write(jsonStr) print( xml_json finished! ) print(len(xml_dir))pythonXmlToJson(path)

02.generate_labels.py

# coding:utf-8import osimport os.path as pathimport json# 保存所有类别对应的id的字典class_id_dict = {}# 判断是不是图片def isPic(basename): file_type = basename.split( . )[-1] pic_file_list = [ png , jpg , jpeg , BMP , JPEG , JPG , JPeG , Jpeg , PNG , TIF , bmp , tif ] if file_type in pic_file_list: return True return False# 判断这个图片有没有对应的json文件def has_json(img_file): # 得到json文件名 base_name = path.basename(img_file) dir_name = img_file[:len(img_file) - len(base_name)] json_name = base_name.split( . )[0] json_name = json_name + .json  json_name = path.join(dir_name, json_name)  if path.isfile(json_name): return json_name return None# 生成file_name的label文件，并重新写入 content_list 中内容def rewrite_labels_file(file_name, content_list): with open(file_name, w ) as f: for line in content_list: curr_line_str =  for element in line: curr_line_str += str(element) +  f.write(curr_line_str + \\n ) return# 生成file_name的训练图片路径文件def rewrite_train_name_file(file_name, content_list): with open(file_name, w ) as f: for line in content_list: f.write(str(line) + \\n ) return # 读取文件def read_file(file_name): if not path.exists(file_name): print( warning:不存在文件 +str(file_name)) return None with open(file_name, r , encoding= utf-8 ) as f: result = [] for line in f.readlines(): result.append(line.strip( \\n )) return result# 加载class_iddef load_class_id(class_name_file): global class_id_dict class_list = read_file(class_name_file) for i in range(len(class_list)): class_id_dict[str(class_list[i])] = i return class_id_dict# 得到分类的id,未分类是-1def get_id(class_name, class_name_file): global class_id_dict if len(class_id_dict) 1: class_id_dict = load_class_id(class_name_file) print( 分类 id 加载完成 ) class_name = class_name if class_name in class_id_dict.keys(): return class_id_dict[class_name] return -1# 解析一个points,得到坐标序列def get_relative_point(img_width, img_height, point_list): # point_list是一个包含两个坐标的list dh = 1.0/ img_height x_min = min(point_list[0][0], point_list[1][0]) y_min = min(point_list[0][1], point_list[1][1]) x_max = max(point_list[0][0], point_list[1][0]) y_max = max(point_list[0][1], point_list[1][1]) dw = 1.0 / img_width dh = 1.0/ img_height # 中心坐标 x = (x_min + x_max)/2.0 y = (y_min + y_max)/2.0 w = x_max - x_min h = y_max - y_min x = x*dw w = w*dw y = y*dh h = h*dh return [x, y, w, h]# 解析json文件def paras_json(json_file, class_name_file): if not path.exists(json_file): print( warning:不存在json文件 + str(json_file)) assert(0) f = open(json_file) setting = json.loads(f.read()) # print(setting) # f.close() shapes = setting[ annotation ]  width = float(setting[ annotation ][ size ][ width ])  height = float(setting[ annotation ][ size ][ height ])  # 拿到标签坐标 result = [] flag = 0 count = 0 if object not in [x for x in shapes]: return [[0,0,0,0,0]] for shape in shapes[ object ]: count += 1 a = [x for x in shapes[ object ]] # print(shape[ name ]) if count == 5 and a[0] == name : flag = 1 else: flag = count if flag == 1: shape = shapes[ object ] class_name = shape[ name ] #得到分类名 class_id = get_id(class_name, class_name_file) point = [] a = [] b = [] a.append(float(shape[ bndbox ][ xmin ])) a.append(float(shape[ bndbox ][ ymin ])) b.append(float(shape[ bndbox ][ xmax ])) b.append(float(shape[ bndbox ][ ymax ])) point.append(a) point.append(b) print(point) locate_result = get_relative_point(width, height, point) locate_result.insert(0, class_id) result.append(locate_result) return result if flag == count: for shape in shapes[ object ]: class_name = shape[ name ] #得到分类名 class_id = get_id(class_name, class_name_file) point = [] a = [] b = [] a.append(float(shape[ bndbox ][ xmin ])) a.append(float(shape[ bndbox ][ ymin ])) b.append(float(shape[ bndbox ][ xmax ])) b.append(float(shape[ bndbox ][ ymax ])) point.append(a) point.append(b) print(point) locate_result = get_relative_point(width, height, point) # 插入id locate_result.insert(0, class_id) result.append(locate_result) print(result) return result# 得到文件夹下所有的图片文件def get_pic_file_from_dir(dir_name): return:所有的图片文件名 if not path.isdir(dir_name): print( warning:路径 %s 不是文件夹 %dir_name) return [] result = [] for f in os.listdir(dir_name): curr_file = path.join(dir_name, f) if not path.isfile(curr_file): continue if not isPic(curr_file): continue result.append(f) return resultdef main(class_name= classes.names , img_dir= dms/images/ , train_txt= dms/train.txt , labels_dir= dms/labels ): cwd = os.getcwd() img_dir = path.join(cwd, img_dir) # print(img_dir) labels_dir = path.join(cwd, labels_dir) if not path.exists(img_dir): print( error:没有发现图片文件夹 , img_dir) if not path.exists(labels_dir): os.mkdir(labels_dir) count = 0  dir_len = len(os.listdir(img_dir)) # 进度条 # print(dir_len) imgs = [] for f in os.listdir(img_dir):  # print(f) curr_path = path.join(img_dir, f) # print(curr_path) if not path.isdir(curr_path): # 不是文件夹就先跳过 continue curr_train_dir = curr_path # print(curr_train_dir) # 是文件夹就创建labels对应的文件夹 curr_labels_dir = path.join(labels_dir, f) if not path.isdir(curr_labels_dir): os.mkdir(curr_labels_dir) # 拿到文件夹下所有的图片文件 curr_dir_imgs = get_pic_file_from_dir(curr_train_dir) # print(curr_dir_imgs) # 解析这些图片的json文件 for img_file in curr_dir_imgs: curr_img_file = path.join(curr_train_dir, img_file) # print(curr_img_file) json_file = has_json(curr_img_file) print(json_file) if json_file: # 保存图片路径 imgs.append(curr_img_file) # 得到json信息 list json_inf = paras_json(json_file, class_name) # print(json_inf) # 标签文件名 label_name = img_file.split( / )[-1].split( . )[0] + .txt  curr_labels_file = path.join(curr_labels_dir, label_name) # 写入标签 rewrite_labels_file(curr_labels_file, json_inf) count += 1 print( \\r当前进度: {:02f} % .format(count/dir_len * 100.0), end= ) print( \\n 保存训练图片路径到: , train_txt) rewrite_train_name_file(train_txt, imgs) return if __name__ == __main__ : main()

03.k_means.py

# coding: utf-8# 获得anchorsfrom __future__ import division, print_functionimport osimport os.path as pathimport jsonimport mathimport numpy as npimport generate_labels as tool# 生成file_name文件，并追加写入name_list中内容def write_file(file_name, name_list): base_name = path.basename(file_name) dir_name = file_name[:len(file_name)-len(base_name)] if not path.exists(dir_name): os.mkdir(dir_name) with open(file_name, a ) as f: for name in name_list: f.write(name+ \\n ) return # 生成file_name文件，并重新写入name_list中内容def rewrite_file(file_name, name_list): base_name = path.basename(file_name) dir_name = file_name[:len(file_name)-len(base_name)] if not path.exists(dir_name): os.mkdir(dir_name) with open(file_name, w ) as f: for i in range(len(name_list)): f.write(str(i) + + name_list[i] + \\n ) return # 解析一个points,得到坐标序列def get_point(point_list): # point_list是一个包含两个坐标的list x_min = min(point_list[0][0], point_list[1][0]) y_min = min(point_list[0][1], point_list[1][1]) x_max = max(point_list[0][0], point_list[1][0]) y_max = max(point_list[0][1], point_list[1][1]) result = str(x_min) + + str(y_min) + + str(x_max) + + str(y_max) return result# 解析json文件def paras_json(json_file, class_name_file): if not path.exists(json_file): print( warning:不存在json文件 + str(json_file)) assert(0) f = open(json_file) setting = json.loads(f.read()) # print(setting) # f.close() shapes = setting[ annotation ]  height = setting[ annotation ][ size ][ height ] width = setting[ annotation ][ size ][ width ] # 拿到标签坐标 result =  flag = 0 count = 0 if object not in [x for x in shapes]: result += + 0 + + get_point([[0,0],[0,0]]) return str(width) + + str(height) + result for shape in shapes[ object ]: count += 1 a = [x for x in shapes[ object ]] # print(shape[ name ]) if count == 5 and a[0] == name : flag = 1 else: flag = count if flag == 1: shape = shapes[ object ] class_name = shape[ name ] #得到分类名 class_id = tool.get_id(class_name, class_name_file) point = [] a = [] b = [] a.append(float(shape[ bndbox ][ xmin ])) a.append(float(shape[ bndbox ][ ymin ]))  b.append(float(shape[ bndbox ][ xmax ])) b.append(float(shape[ bndbox ][ ymax ])) point.append(a) point.append(b) # print(point) locate_result = get_point(point) result += + str(class_id) + + locate_result return str(width) + + str(height) + result if flag == count: for shape in shapes[ object ]: class_name = shape[ name ] #得到分类名 class_id = tool.get_id(class_name, class_name_file) point = [] a = [] b = [] a.append(float(shape[ bndbox ][ xmin ])) a.append(float(shape[ bndbox ][ ymin ]))  b.append(float(shape[ bndbox ][ xmax ])) b.append(float(shape[ bndbox ][ ymax ])) point.append(a) point.append(b) # print(point) locate_result = get_point(point) result += + str(class_id) + + locate_result return str(width) + + str(height) + result# 得到文件夹里面所有的图片路径def get_pic_file(dir_path): result = [] if not path.isdir(dir_path): print( exception: 路径%s不是文件夹 %(dir_path)) return result # 读取图片路径 for f in os.listdir(dir_path): curr_file = path.join(dir_path, f) if not path.isfile(curr_file): continue if not tool.isPic(f): continue result.append(curr_file) return result# 生成k_means需要的数据def generate_k_means_data(class_name= train.names , train_dir= ./JPEGImages ): train_list = [] img_index = 0 dir_list = os.listdir(train_dir) for i in range(len(dir_list)): f = dir_list[i] curr_path = os.path.join(train_dir, f) if not path.isdir(curr_path): continue curr_dir_imgs = get_pic_file(curr_path) # 判断有没有对应的json文件,并解析json文件保存到list中 for img_file in curr_dir_imgs: json_file = tool.has_json(img_file) if json_file: # 有这个json文件就保存这个img json_inf = paras_json(json_file, class_name) train_list.append(str(img_index) + + str(img_file) + + str(json_inf)) img_index += 1 print( \\r文件夹处理进度：{:2f}% .format( (i+1)*100 / dir_list.__len__() ), end= ) print() return train_list