一、标注数据
使用labelImg进行标注,标注方法自行研究,标注完成后生成的是xml后缀的文件,该格式是VOC的标准格式。
二、制作自己的数据集
我这里总共标注训练数据4351张,类别为船(共一类),数据集制作参考yoloV4里面的readme。
import os
import xml.etree.ElementTree as ET
if __name__ == '__main__':
list_file = os.listdir('F:\\155dataSet\\boatAll\\annotations\\')
for i in list_file:
annotation_file = f"F:\\155dataSet\\boatAll\\annotations\\{i}"
txtfile = f"F:\\155dataSet\\boatAll\\labels\\{i}"
txt_file = f"{txtfile[0:-4]}.txt"
objects = ET.parse(annotation_file).findall("object")
sizes = ET.parse(annotation_file).findall("size")
for size in sizes:
width = float(size.find('width').text)
height = float(size.find('height').text)
for object in objects:
bbox = object.find('bndbox')
xmin = float(bbox.find('xmin').text) - 1
ymin = float(bbox.find('ymin').text) - 1
xmax = float(bbox.find('xmax').text) - 1
ymax = float(bbox.find('ymax').text) - 1
classname = object.find('name').text
xcenter = (xmax + xmin) / (2 * width)
ycenter = (ymax + ymin) / (2 * height)
w = (xmax - xmin) / (width)
h = (ymax - ymin) / (height)
if classname == 'boat':
object_class = 0
else:
object_class = 1
data = [str(object_class),' ',str(xcenter),' ',str(ycenter),' ',str(w),' ',str(h)]
with open(txt_file,'a+') as f:
f.writelines(data)
f.write('\n')
该部分代码适用于xml转化为txt文件,文件格式如下图
生成的txt文件和jpg文件放到同一个目录下,这里我放到了./build/darknet/x64/data/obj。
接下来,就是按照readme的操作制作cfg,train.txt,obj.data.obj.names文件。
2.1. 创建文件yolo-obj.cfg,修改配置文件
点1 三个yolo层前面的卷积层,修改方法为filters=(classes + 5)x3),例如,我就一类,那么filters为18=(1+5)*3
2.2 yolo层classes改为1
2.3 可以根据自己的情况修改batch等一些训练参数
2.4 在build\darknet\x64\data\内创建obj.names
2.5 在build\darknet\x64\data\内创建obj.data
2.6 在build\darknet\x64\data\内创建train.txt
三、训练
你可以从头开始训练,也可以使用预训练的模型
./darknet detector train data/obj.data yolo-obj.cfg yolov4.conv.137
四、剪枝
剪枝操作我研究不深,完全按照github中的yolov3-channel-and-layer-pruning,使用readme中的稀疏和剪枝操作,对训练好的模型进行剪枝,得到best.pt文件,具体的剪枝方法可以参考readme,里面非常详细。
python slim_prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --global_percent 0.8 --layer_keep 0.01
五、微调训练
我先将pt格式的文件转化为了weights类型的文件,该文件比较容易在darknet中进行微调训练,即我将该weights文件和对应的cfg文件放到darknet中,重新微调训练,训练时可以调整一下参数和训练次数等。
python3 -c “from models import *; convert(‘cfg/prune_0.8_keep_0.01_yolov4.cfg’, ‘weights/best.pt’)”
六、测试