最近打算记录一些自己常用的脚本文件,方便之后需要的时候能够查看
1 根据一定的比例划分数据集
在数据集文件中,我们需要将按照一定整个数据集划分成训练集( training dataset) 和测试集 (test dataset)。darknet的数据集目录结构是这样的:
darknet-master
---- data
---- obj.names # 物体类别名称(如果有两类物体,就写上两类物体的名称)
---- obj.data # 将数据集的信息保存在这个文件中,yolov4从这个文件中读取数据集信息
---- obj # 存放图片以及每个图片的标签信息
---- train.txt # 存放训练集地址 (相对地址,比如: data/obj/image1.jpg)
---- test.txt # 存在训练集地址 (相对地址,比如: data/obj/image3.jpg)
我们将在脚本执行的文件目录下新建一个data文件夹,将训练样本和测试样本的地址分别保存data/train.txt和data/test.txt文件中。脚本的内容如下 (脚本名称是 partitioin_dataset.py):
import os
import random
import shutil
import argparse
# step3: 将训练集数据和测试集的图片地址写在训练集和测试集txt文件中
# step3: 将训练集数据和测试集的图片地址写在训练集和测试集txt文件中
def write_path_in_txt(datasets_path,txt_path):
# 如果文件不存在,就创建一个文件
with open(txt_path,'a+') as f:
f.truncate(0)
for img_path in datasets_path:
# 将xml后缀改成jpg
img_path = img_path.replace("xml","jpg")
info = 'data/obj/{}\n'.format(img_path)
#print(info)
f.write(info)
# step2. 划分数据集
# path_dataset: 数据集的地址
# ratio: 训练集所占的比重
def partition_dataset(path_dataset,ratio = 0.8):
if ratio < 0.5 or ratio > 1:
print("0.5 < ratio < 1")
return
# 得到所有的xml文件,将其放在datasets_xml文件中
for _,_,imgs_path in os.walk(path_dataset):
datasets_xml = [img for img in imgs_path if img.endswith('xml')]
# 得到训练数据集和测试数据集
train_list, test_list = [], []
train_list = random.sample(datasets_xml,int(len(datasets_xml)*ratio))
# 遍历整个数据集,将不在训练集中的图片写到到验证集中
for names in datasets_xml:
if names not in train_list:
test_list.append(names)
# 将训练集和测试集写到train.txt和text.txt文件夹中
write_path_in_txt(train_list,'./data/train.txt')
write_path_in_txt(test_list,'./data/test.txt')
def get_parser():
parser = argparse.ArgumentParser(description="partition dataset by ratio")
parser.add_argument('--path', help="the absolute path of the dataset")
parser.add_argument('--ratio', type=float, default = 0.8, help="the ratio of the training set, 0.5 < ratio < 1")
return parser
if __name__ == '__main__':
parser = get_parser()
args = parser.parse_args()
path = args.path
ratio = args.ratio
# step1. 首先创建一个文件夹(如果存在了,就不创建了)
if not os.path.exists('./data'):
os.makedirs('./data')
partition_dataset(path,ratio)
###################################### 使用样例 ####################################
# python partitioin_dataset.py --path C:\Users\cumt\Desktop\path_dataset --ratio 0.8
2 改变标签的坐标
使用labelImage标注的坐标是 (x1, y2, x2, y2),但是darknet需要的坐标形式如下
<object-class> <x_center> <y_center> <width> <height>
object-class: 表示物体的数字标签。比如0, 1, 2等整数。
<x_center> <y_center> <width> <height>: 表示物体的相对中心坐标及其相对宽高(这四个值的大小在0-1之间)。
举例来说:
<x_center> = <absolute_x> / <image_width> = bounding box中心x实际坐标 / 图片实际宽度
<y_center> = <absolute_y> / <image_height> = bounding box中心y实际坐标 / 图片实际高度
<width> = <absolute_width> / <image_width> = bbox宽度 / 图片实际宽度
<height> = <absolute_width> / <image_width> = bbox高度 / 图片实际高度
如果一幅图片中包含不止一个物体,那么每幅图片的标签信息应该如何填写?举例来说,对于image1.jpg图片有三个物体,image1.txt文件就是这样:
1 0.716797 0.395833 0.216406 0.147222
0 0.687109 0.379167 0.255469 0.158333
1 0.420312 0.395833 0.140625 0.166667
具体内容参考我的这一篇博客。
我们会遍历数据集下所有的xml文件,然后读取信息并将转换之后的坐标信息,以txt的形式保存在数据集的同目录下。最终的代码就是这样的:
'''
这个脚本文件用来将(x1,y1,x2,y2)的坐标转成darkent-yolov4的形式
读取文件中所有的xml文件,然后生成对应txt文件,txt文件在数据集的同目录下
'''
import xml.etree.ElementTree as ET
from tqdm import tqdm
import argparse
import os
def transformation(path_dataset,obj_names_lst):
# 使用labelImage标注图片的时候,有的xml文件中关于图片的width和height可能为0
empty_list = []
for path_root, _ ,path_files_list in os.walk(path_dataset):
for path_file in tqdm(path_files_list,total=len(path_files_list),unit='fils'):
if path_file.endswith('xml'):
name = path_file
path_file = os.path.join(path_root,path_file)
# 读取xml文件,然后将其写入txt文件中
# 解析xml文件
tree = ET.parse(path_file)
# 得到根结点
root = tree.getroot()
# 获取图片的width和height
for sizeNode in root.iter('size'):
width_img = float(sizeNode.find('width').text)
height_img = float(sizeNode.find('height').text)
# 打开txt文件,将内容写进去
path_txt = os.path.splitext(path_file)[0]+'.txt'
# path_list.append(path_txt)
with open(path_txt,'a+') as f:
# 首先先清空txt文件
f.truncate(0)
# 打开xml文件,name, xmin,ymin,xmax,ymax信息
for objectNode in root.iter('object'):
cls_name = objectNode.find('name').text
for index, name in enumerate(obj_names_lst):
if cls_name == obj_names_lst[index]:
obj_cls = str(index)
break
# # 获取obj类别标签
# if cls_name == 'person':
# obj_cls = str(0)
# elif cls_name == 'hat':
# obj_cls = str(1)
# 得到xmin,ymin,xmax,ymax
xmin = float(objectNode.getchildren()[4].find('xmin').text)
ymin = float(objectNode.getchildren()[4].find('ymin').text)
xmax = float(objectNode.getchildren()[4].find('xmax').text)
ymax = float(objectNode.getchildren()[4].find('ymax').text)
# 得到x_center,y_center,width,height
# x_center = (xmin+xmax)/2; y_center = (ymin_ymax) / 2; width = (xmax-xmin); height = (ymax-ymin)
x_center, y_center, width, height = (xmin+xmax) / 2, (ymin+ymax) / 2, xmax-xmin, ymax-ymin
# 得到相对的值
# 有的值为零
if (0 ==width_img or 0 == width_img):
if not name in empty_list:
empty_list.append(name)
else:
abs_x = float(x_center / width_img)
abs_y = float(y_center / height_img)
abs_width = float(width / width_img)
abs_height = float(height / height_img)
info = obj_cls +" " + str(abs_x) + " " + str(abs_y) + " " + str(abs_width) + " " + str(abs_height) + "\n"
# 创建相关的文件,将内容写进去
f.write(info)
return empty_list
def get_parser():
parser = argparse.ArgumentParser("coordination transformation: (x1,y1,x2,y2) -> (abs_x, abs_y, abs_w, abs_h)")
parser.add_argument("--path",help="the absolute path of the dataset", default=str)
parser.add_argument("--obj_names",help="obj names",nargs="+")
return parser
if __name__ == '__main__':
parser = get_parser()
args = parser.parse_args()
path = args.path
obj_names_lst = args.obj_names
empty_list = transformation(path,obj_names_lst)
if (len(empty_list)):
print("##################### the width and height of these xml are zero #####################")
print(empty_list)
###################################### 使用样例 ####################################
# --obj_names输入的是从0开始的标签值,使用的时候就是一个列表
# python coord_transformation.py --path C:\Users\cumt\Desktop\path_dataset --obj_names person hat
3 计算数据集中每类bounding box的个数
最近在整理一个数据集,需要知道该数据集中不同类别的bbox的个数,所以自己就写了一个脚本文件 (脚本名称是get_num_bboxes.py),这个脚本文件会读取数据集下所有的xml文件,读取其中的name,最终以词典的形式输出打印。
import os
import argparse
from tqdm import tqdm
import xml.etree.ElementTree as ET
def calc_bbox_num(input_path):
# 判断是是否是一个路径
if (not os.path.isdir(input_path)):
print("input path is not a folder path")
names_bbox = {}
for path_root, _, name_file_list in os.walk(input_path):
for path_file in tqdm(name_file_list,total=len(name_file_list),unit='xml file'):
if path_file.endswith('xml'):
# 读取xml文件 (要得到绝对路径)
path_file = os.path.join(path_root,path_file)
tree = ET.parse(path_file)
root = tree.getroot()
# 得到所有的object结点
for objectNode in root.iter('object'):
cls_name = objectNode.find('name').text
if cls_name not in names_bbox.keys():
names_bbox[cls_name] = 1
else:
names_bbox[cls_name] += 1
print(names_bbox)
def get_parser():
parser = argparse.ArgumentParser(description="calculate bounding box's number of one dataset ")
# 输入的是一个绝对路径
parser.add_argument('--path',help='the absolute path of the dataset')
return parser
if __name__ == '__main__':
parser = get_parser()
args = parser.parse_args()
input_path = args.path
calc_bbox_num(input_path)
###################################### 使用样例 ####################################
# python get_num_bboxes.py --path C:\Users\cumt\Desktop\path_dataset
最终的运行结果
读取的过程会有进度条进行动态显示