voc装yolo|划分训练集和验证集|yolo格式标签画在图上保存||查看yolo格式数据集详细信息-CSDN博客

本文链接：https://blog.csdn.net/qq_15060477/article/details/141754834
import xml.etree.ElementTree as ET
import os
import argparse
import random
import shutil
import cv2
rootpath = os.getcwd()

class VocToYolo(object):
    '''
    labelimg贴的标签voc文件转化为yolo支持的格式
    '''
    def __init__(self,xml_path,save_txt_path):
        self.xml_path =xml_path
        self.save_txt_path = save_txt_path
        self.roots = self.readxmls()
    def readxmls(self):
        f = open(self.xml_path)
        xml_text = f.read()
        root = ET.fromstring(xml_text)
        f.close()
        return root
    def get_image_path(self):
        root = self.readxmls()
        images_path = root.find('filename')
        img_name = images_path.text.split('/')[-1].split('.')[0]
        print(img_name)
        return images_path.text,img_name
    def get_hw(self):
        size = self.roots.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        return (w,h)
    def convert_size(self,size, box):
        dw = 1.0 / size[0]
        dh = 1.0 / size[1]
        x = (box[0] + box[1]) / 2.0
        y = (box[2] + box[3]) / 2.0
        w = box[1] - box[0]
        h = box[3] - box[2]
        x = x * dw
        w = w * dw
        y = y * dh
        h = h * dh
        return (x, y, w, h)
    def get_xyxy_label(self,classes):
        _,a = self.get_image_path()
        os.makedirs(self.save_txt_path,exist_ok=True)
        out_file = open(self.save_txt_path+'/' + str(a) + '.txt', 'w')
        for obj in self.roots.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                print(cls)
                continue
            print("标签：",cls)
            label_now = classes.index(cls)
            print("标签索引:", label_now)
            xmlbox = obj.find('bndbox')
            xmin = int(float(xmlbox.find('xmin').text))
            xmax = int(float(xmlbox.find('ymin').text))
            ymin = int(float(xmlbox.find('xmax').text))
            ymax = int(float(xmlbox.find('ymax').text))
            size = self.get_hw()
            xyhw = self.convert_size(size, (xmin, ymin, xmax, ymax))

            out_file.write(str(label_now) + " " + " ".join([str(xyhw) for xyhw in xyhw]) + '\n')
            print(xmin, ymin, xmax, ymax)
        return 0


class SplitValTrain(object):
    '''
    划分验证集和训练集
    '''
    def __init__(self,rootpath,label_path,images_path,ratios,base_save_path):
        self.rootpath = rootpath
        self.label_path = label_path
        self.images_path = images_path
        self.ratios = ratios
        # self.train_images_path = rootpath + '/yolodata/train/images'
        # self.train_labels_path = rootpath + '/yolodata/train/labels'
        # self.val_images_path = rootpath + '/yolodata/val/images'
        # self.val_labels_path = rootpath + '/yolodata/val/labels'


        self.train_images_path = rootpath +'/'+base_save_path+ '/images/train'
        self.val_images_path = rootpath +'/'+base_save_path+  '/images/val'

        self.train_labels_path = rootpath +'/'+base_save_path+  '/labels/train'
        self.val_labels_path = rootpath + '/'+base_save_path+ '/labels/val'
    def make_data_dir(self):
        os.makedirs(self.train_images_path,exist_ok=True)
        os.makedirs(self.train_labels_path,exist_ok=True)
        os.makedirs(self.val_images_path, exist_ok=True)
        os.makedirs(self.val_labels_path, exist_ok=True)
    def split_val_train(self):
        self.make_data_dir()
        filename = os.listdir(self.images_path)
        # print(filename)
        random.shuffle(filename)

        go_to = int(self.ratios*len(filename))
        print(go_to,len(filename))
        for i in range(len(filename)):
            name_num = filename[i].split('.')[0]
            print(filename[i])
            try:
                if i<go_to:
                    shutil.copy(self.images_path+'/'+str(filename[i]), self.train_images_path)
                    shutil.copy(self.label_path + '/' + str(name_num) + '.txt', self.train_labels_path)
                else:
                    shutil.copy(self.images_path + '/' + str(filename[i]), self.val_images_path)
                    shutil.copy(self.label_path + '/' + str(name_num) + '.txt', self.val_labels_path)
            except:
                continue

        return 0
def plot_images(image_path,txt_path,opt):
    # 加载图像
    image = cv2.imread(image_path)

    # 假设txt文件中的内容格式如下：
    # class_id x_center y_center width height
    # 其中x_center, y_center, width, height是相对于图像的比例值，范围从0到1
    with open(txt_path, 'r') as f:
        lines = f.readlines()

    # 画框和标签
    for line in lines:
        class_id, x_center, y_center, width, height = map(float, line.split())

        # 将比例坐标转换为像素坐标
        top = int(image.shape[0] * y_center - image.shape[0] * height / 2)
        bottom = int(image.shape[0] * y_center + image.shape[0] * height / 2)
        left = int(image.shape[1] * x_center - image.shape[1] * width / 2)
        right = int(image.shape[1] * x_center + image.shape[1] * width / 2)

        # 确保坐标不会超出图像边界
        top = max(0, top)
        left = max(0, left)
        bottom = min(image.shape[0], bottom)
        right = min(image.shape[1], right)

        # 画框
        cv2.rectangle(image, (left, top), (right, bottom), opt.colors, opt.font_thickness)

        # 标签（可选）
        # label = f"Class {stuff_names[int(class_id)]}"
        label = f" {opt.stuff_names[int(class_id)]}"
        cv2.putText(image, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, opt.font_size, opt.colors, opt.font_thickness)

    # 显示图像
    # cv2.imshow('Image with Bounding Boxes', image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    return image,int(class_id)
def plotimgmore(opt):

    # base_txt_path = 'road_large_yolodata/labels/train'
    # base_img_path = 'road_large_yolodata/images/train'
    # save_img_path = './result'
    os.makedirs(opt.save_img_path,exist_ok=True)
    txt_list = os.listdir(opt.base_txt_path)
    for txt_file_name in txt_list:
        try:
            txt_path = os.path.join(opt.base_txt_path,txt_file_name)
            img_path = os.path.join(opt.base_img_path,txt_file_name.split('.')[0]+opt.img_fomate_name)
            # img_path =txt_path.replace('txt',opt.img_fomate_name).replace('labels','images')
            print('txt_path::: ',txt_path)
            print('img_path::: ', img_path)
            res,index = plot_images(img_path,txt_path,opt)

            # if stuff_names[index] in acounter_list:
            cv2.imwrite(os.path.join(opt.save_img_path,txt_file_name.split('.')[0]+'.jpg'),res)
            # del acounter_list[index]
        except:
            print('hhhh')

def excute_voc_to_yolo():
    '''
    转化为yolo格式主函数，会产生一个文件夹，文件夹中含有yolo格式的txt文件
    '''
    classes = ["with_mask", "without_mask"]
    parser = argparse.ArgumentParser()
    parser.add_argument('--xml_path', type=str, default='./mask_detect/annotations', help='xml路径')
    parser.add_argument('--save_txt', type=str, default='./labels', help='txt保存路径')
    parser.add_argument('--classes', type=list, default=["with_mask", "without_mask"], help='xml中类别名称类别名')
    opt = parser.parse_args()
    ba = opt.classes
    for i in os.listdir(opt.xml_path):
        xml_path = opt.xml_path + '/' + str(i)
        ab = VocToYolo(xml_path, opt.save_txt)
        ab.get_xyxy_label(opt.classes)
def excute_split_val_train():
    '''划分验证集和训练集'''
    parser = argparse.ArgumentParser()
    parser.add_argument('--rootpath', type=str, default=rootpath, help='当前目录')
    parser.add_argument('--label_path', type=str, default='./yolo_labels', help='txt文件路径')
    parser.add_argument('--images_path', type=str, default='./images/images', help='图像路径')
    parser.add_argument('--ratios', type=float, default=0.8, help='训练集划分比例')
    parser.add_argument('--base_save_path', type=str, default='self_driving_5class', help='划分后数据集的主文件夹名')
    opt = parser.parse_args()
    ab = SplitValTrain(opt.rootpath, opt.label_path, opt.images_path, opt.ratios,opt.base_save_path)
    ab.split_val_train()

def plot_label_to_img():
    '''
    将标签画在图像上，并保存
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument('--base_txt_path', type=str, default='./yolo_labels', help='标签文件路径')
    parser.add_argument('--base_img_path', type=str, default='./images/images', help='标签文件所对应的图像模路径')
    parser.add_argument('--save_img_path', type=str, default='./plotreslts', help='图像存储路径')
    parser.add_argument('--stuff_names', type=list, default= ['car','truck','person','bicycle','traffic_light'] ,help='要画到图像上的类别标签')
    parser.add_argument('--img_fomate_name', type=str, default='.jpg', help='图像格式后缀名')
    parser.add_argument('--font_size', type=int, default=0.8, help='绘制字体大小')
    parser.add_argument('--font_thickness', type=int, default=2, help='字体粗细')
    parser.add_argument('--colors', type=tuple, default=(0,255,0), help='字体颜色')


    opt = parser.parse_args()
    plotimgmore(opt)
def create_text_for_for_data():
    stuff_names =  ['car','truck','person','bicycle','traffic_light']
    class_number = len(stuff_names)
    data_name = '数据集类型'
    tarin_txt = './yolodata/labels/train'
    val_txt = './yolodata/labels/val'
    test_txt = './yolodata/labels/test'
    train_number=len(os.listdir(tarin_txt))
    val_number=len(os.listdir(val_txt))
    try:
        total_data_number=len(os.listdir(tarin_txt))+len(os.listdir(val_txt))+len(os.listdir(test_txt))
        print("yolo格式数据集|%s|%s类别|" % (
        data_name, class_number))
        print("yolo格式数据集|%s|%s类别" % (data_name, class_number))
        print(
            "本数据为%s检测数据集，数据集数量如下：\n总共有:%s张\n训练集：%s张\n验证集：%s张\n测试集：%s\n类别数量：%s\n类别名：\n%s" % (
            data_name, total_data_number,
            train_number, val_number, len(os.listdir(test_txt)), class_number, stuff_names))
    except:
        total_data_number = len(os.listdir(tarin_txt)) + len(os.listdir(val_txt))
        print("yolo格式数据集|%s|%s类别" % (data_name, class_number))
        print("olo格式数据集|%s|%s类别|"%(data_name,class_number))
        print("本数据为%s检测数据集，数据集数量如下：\n总共有:%s张\n训练集：%s张\n验证集：%s张\n类别数量：%s\n类别名：\n%s"%(data_name,total_data_number,
                                            train_number,val_number,class_number, stuff_names))
        print("无测试集")
if __name__ == '__main__':
    #将xml格式的数据集转化为yolo格式的数据集，使用时只需要修改如下函数里面的参数即可自动产生主文件夹名字为labels的txt文件
    # excute_voc_to_yolo()
    # 划分训练集和验证集，使用时只需要修改如下函数里面的参数即可
    # excute_split_val_train()
    # 将标签画在图像上，并保存
    # plot_label_to_img()
    # 查看数据集信息
    create_text_for_for_data()