1. 将xml格式的数据转换为yolo格式
由于voc对数据集的标注格式是xml文件的，标注信息清楚；而对于yolo的标注格式是txt文件的，由目标类别和标注框的相对位置来组成：
# voc的xml文件内容，以2007_000027.xml为例(省略部分内容)： <annotation> <folder>VOC2012</folder> <filename>2007_000027.jpg</filename> <source> <database>The VOC2007 Database</database> <annotation>PASCAL VOC2007</annotation> <image>flickr</image> </source> <size> <width>486</width> <height>500</height> <depth>3</depth> </size> <segmented>0</segmented> <object> <name>person</name> <pose>Unspecified</pose> <truncated>0</truncated> <difficult>0</difficult> <bndbox> <xmin>174</xmin> <ymin>101</ymin> <xmax>349</xmax> <ymax>351</ymax> ...... </object> </annotation> # yolo需要的标注格式： 14 0.5380658436213992 0.452 0.360082304526749 0.5
这里直接将xml目录里的文件全部转换成txt文件，并存储在所设定目录下
- 代码如下：(xml_to_txt.py)
import os import pickle import xml.etree.ElementTree as ET from tqdm import tqdm # chance your classes here names: ['锅', '陶瓷器皿', '鞋', '食用油桶', '饮料瓶', '鱼骨', '一次性快餐盒', '书籍纸张', '充电宝', '剩饭剩菜', '包', '垃圾桶', '塑料器皿', '塑料玩具', '塑料衣架', '大骨头', '干电池', '快递纸袋', '玻璃器皿', '砧板', '筷子', '纸盒纸箱', '花盆', '茶叶渣', '菜帮菜叶', '蛋壳', '调料瓶', '软膏', '过期药物', '酒瓶', '金属厨具', '金属器皿', '金属食品罐', '插头电线', '旧衣服', '易拉罐', '枕头', '果皮果肉', '毛绒玩具', '污损塑料', '污损用纸', '洗护用品', '烟蒂', '牙签'] # function: # (xmin, xmax, ymin, ymax) -> (center_x, cneter_y, box_weight, box_height) # size: (w, h) # box : (xmin, xmax, ymin, ymax) def convert(size, box): dw = 1.0 / size[0] dh = 1.0 / size[1] x = (box[0] + box[1]) / 2.0 # center_x y = (box[2] + box[3]) / 2.0 # center_y w = box[1] - box[0] # box_weight h = box[3] - box[2] # box_height x = x * dw w = w * dw y = y * dh h = h * dh return (x, y, w, h) # 功能：对单个文件进行转换，其中名字是相匹配的 # infile: xml文件路径 # outfile：txt文件输出路径 def convert_annotation(infile, outfile): in_file = open(infile, encoding='utf-8') # xml path out_file = open(outfile, '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) 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) xmlbox = obj.find('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') in_file.close() out_file.close() # 功能：对xml目录下的全部xml文件全部转换成yolo的txt格式，并输出在savepath目下 # xmlpath: the dir path save xml files # savepath: the dir path to save txt files def xml_to_txt(xmlpath, savepath): assert os.path.exists(xmlpath), "{} not exists.".format(xmlpath) if not os.path.exists(savepath): os.mkdir(savepath) assert os.path.exists(savepath), "{} create fail.".format(savepath) xmlfiles = os.listdir(xmlpath) assert len(xmlfiles) != 0, "{} have no files".format(xmlpath) xmlfilepaths = [os.path.join(xmlpath, xmlfile) for xmlfile in xmlfiles] for xmlfilepath in tqdm(xmlfilepaths, desc="change xml to txt: "): # E:\学习\机器学习\数据集\众盈比赛数据集\xml\train1.xml -> train1 image_id = xmlfilepath.split('\\')[-1].split('.')[0] # train1 -> \..\train1.txt txtfilepath = os.path.join(savepath, image_id) + ".txt" convert_annotation(xmlfilepath, txtfilepath) if __name__ == '__main__': xmlpath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\Annotations" savepath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\YoloLabels" xml_to_txt(xmlpath, savepath)
- 测试例程：
- 结束后效果：
2. 将数据集切分为训练集与验证集
对于yolo的数据集，所需要切分成的格式为为：
DataSet（file name可修改） |---------images |-------train |-------val |-------test（可以没有） |---------labels |-------train |-------val |-------test（可以没有）
这个固定格式是有VOC.yaml文件所决定的：
也就是在一个数据集下，需要有images与labels两个子文件夹，images内直接用来存放图片，labels用来存放对应图片的标签文件；没有val与test表示训练的过程中只有train的文件。可以看见，只需要设置好train与val的路径即可。所以，要是需要训练自己的数据集时，可以自己写一个脚本文件为数据集来服务。以下就是我自己写的脚本代码。
- 代码如下：(dataset_split.py)
import os from tqdm import tqdm from sklearn.model_selection import train_test_split from shutil import copyfile # 功能: 递归穿件目录文件，有则跳过，没有则创建并验证 def _mkdir(dirpath): if not os.path.exists(dirpath): os.makedirs(dirpath) assert os.path.exists(dirpath), "{} create fail.".format(dirpath) # 功能: 从图片路径中获取最后的文件名(含后缀或者不含后缀) def get_filename(filepath, suffix=True): if suffix is True: image_id = filepath.split('\\')[-1] # 获取文件名(含后缀) else: image_id = filepath.split('\\')[-1].split('.')[0] # 获取文件名(不含后缀) return image_id # 功能: 根据图片数据与标签数据构建yolov5格式的数据集 # imagepath: 存放数据图片的路径 # labelpath: 存放数据标签的路径 # datasetpath: 构建voc格式数据集的目录路径 def dataset_split(imagepath, labelpath, datasetpath): assert os.path.exists(imagepath), "{} not exists.".format(imagepath) assert os.path.exists(labelpath), "{} not exists.".format(labelpath) # 构建图片与标签的路径列表，然后按8:2进行训练与测试的切分 images = os.listdir(imagepath) imagelists = [os.path.join(imagepath, image) for image in images] assert len(imagelists) != 0, "{} have no files".format(imagepath) labels = os.listdir(labelpath) labellists = [os.path.join(labelpath, label) for label in labels] assert len(labellists) != 0, "{} have no files".format(labelpath) X_train, X_test, y_train, y_test = train_test_split(imagelists, labellists, test_size=0.2, random_state=42) # 构建voc格式的数据集存放格式 _mkdir(datasetpath) # 在数据路径datasetpath下，分别递归建立目录: # | -----Dataset # | -----images # | -----train # | -----val # | -----labels # | -----train # | -----val images_train_dirpath = os.path.join(datasetpath, "images", "train") _mkdir(images_train_dirpath) images_val_dirpath = os.path.join(datasetpath, "images", "val") _mkdir(images_val_dirpath) labels_train_dirpath = os.path.join(datasetpath, "labels", "train") _mkdir(labels_train_dirpath) labels_val_dirpath = os.path.join(datasetpath, "labels", "val") _mkdir(labels_val_dirpath) # 复制图片/标签列表中的数据到指定子目录下 for image_path in tqdm(X_train, desc="process train images"): image_newpath = os.path.join(images_train_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) for image_path in tqdm(X_test, desc="process val images"): image_newpath = os.path.join(images_val_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) for image_path in tqdm(y_train, desc="process train labels"): image_newpath = os.path.join(labels_train_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) for image_path in tqdm(y_test, desc="process val labels"): image_newpath = os.path.join(labels_val_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) print("Process Success. The data has save to {}".format(datasetpath)) if __name__ == '__main__': imagepath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\JPEGImages" labelpath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\YoloLabels" datasetpath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\Dataset" dataset_split(imagepath, labelpath, datasetpath)
- 测试例程：
这里我的代码可能没有优化好，数据有点多，感觉处理得也比较慢。
- 结束后效果：
