第一步,数据集:自己截图,不需要resize,使用labelimg标注,标注后分布放入下面文件夹
第二步,运行xml_to_txt.py 代码如下
# _*_ coding: utf-8 _*_
# @Time : 2020/2/28 20:57
# @File : xml_to_txt.py
"""
需要修改的地方 :
1、你自己的类别 CLASSES
2、数据集 image路径 img_dir annotation路径 annotation_dir
3、第22行,生成文件的保存路径及名称
"""
import xml.etree.ElementTree as ET
import os
CLASSES = ['c', 'nc']
def convert_xml_annotation(img_dir,annotation_dir, classes,target_file):
xml_dir = []
for xml in os.listdir(annotation_dir):
if xml.endswith('.xml'):
xml_dir.append(xml)
print("Total xml files : ", len(xml_dir))
with open(target_file, 'w') as f:
for i in range(len(xml_dir)):
tree = ET.parse(os.path.join(annotation_dir,xml_dir[i]))
root = tree.getroot()
# image path
filename = root.find('filename').text
image_path = os.path.join(img_dir,filename)
annotation = image_path
# coordinates of label : xmin ymin xmax ymax
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
bbox = obj.find('bndbox')
xmin = bbox.find('xmin').text.strip()
xmax = bbox.find('xmax').text.strip()
ymin = bbox.find('ymin').text.strip()
ymax = bbox.find('ymax').text.strip()
annotation += ' ' + ','.join([xmin, ymin, xmax, ymax,str(cls_id)])
print(annotation)
f.write(annotation + "\n")
#训练数据
convert_xml_annotation(r'D:\projects\YOLOV3\images\train\image',r"D:\projects\YOLOV3\images\train\annotation", CLASSES,r"D:\projects\YOLOV3\data\dataset\train.txt")
#测试数据
convert_xml_annotation(r'D:\projects\YOLOV3\images\test\image',r"D:\projects\YOLOV3\images\test\annotation", CLASSES,r"D:\projects\YOLOV3\data\dataset\test.txt")
生成文件如下
第三步,新建.names文件,如下
内容如下,我的只有两中类别c,nc
第四步,修改 config.py
配置好,可以开始训练了,我比较懒,只标注了66张图片进行训练,下面是训练后的预测效果,由于训练数据过小,准确率不是很高,也偶尔存在误判
点赞过百上传源码哦~