python常用的脚本文件

最近打算记录一些自己常用的脚本文件，方便之后需要的时候能够查看

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

最终的运行结果

读取的过程会有进度条进行动态显示