VOC转YOLO格式(推荐)

1.xml转换为txt

import os.path
import xml.etree.ElementTree as ET
import os
import random
# class_names = ['palm', 'stone', 'scissor', 'awesome', 'heartB', 'OK', 'ROCK', 'one', 'swear', 'thanks', 'heartA',
#                'heartC', 'good', 'bad', 'pray', 'call', 'take_picture', 'salute']
class_names = ['menopause', 'hairball', 'broken yarn', 'hole','stains']
xmlpath = 'F:/Project_code/yolov7-main/VOCdevkit/VOC2007/Annotations/'  # 原xml路径
txtpath = 'F:/Project_code/yolov7-main/VOCdevkit/VOC2007/labels_copy/'  # 转换后txt文件存放路径
if not os.path.exists(txtpath):
    os.makedirs(txtpath)
files = []

for root, dirs, files in os.walk(xmlpath):
    None

number = len(files)
print(number)
i = 0
while i < number:

    name = files[i][0:-4]
    xml_name = name + ".xml"
    txt_name = name + ".txt"
    xml_file_name = xmlpath + xml_name
    txt_file_name = txtpath + txt_name

    xml_file = open(xml_file_name)
    tree = ET.parse(xml_file)
    root = tree.getroot()
    filename = root.find('filename').text

    image_name = root.find('filename').text
    w = int(root.find('size').find('width').text)
    h = int(root.find('size').find('height').text)

    f_txt = open(txt_file_name, 'w+')
    content = ""

    first = True

    for obj in root.iter('object'):

        name = obj.find('name').text
        class_num = class_names.index(name)

        xmlbox = obj.find('bndbox')

        x1 = int(xmlbox.find('xmin').text)
        x2 = int(xmlbox.find('xmax').text)
        y1 = int(xmlbox.find('ymin').text)
        y2 = int(xmlbox.find('ymax').text)

        if first:
            content += str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)
            first = False
        else:
            content += "\n" + \
                       str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)

    # print(str(i / (number - 1) * 100) + "%\n")
    print(content)
    f_txt.write(content)
    f_txt.close()
    xml_file.close()
    i += 1

print("done!")

2.数据集划分

import os
import random
import cv2

trainval_percent = 0.9
train_percent = 0.95

data_root = "F:/Project_code/yolov7-main/VOCdevkit/VOC2007/"

fdir = data_root + 'ImageSets/Main/'
if not os.path.exists(fdir):
    os.makedirs(fdir)
# xmlfilepath = data_root + 'txts/'
xmlfilepath = data_root + 'Annotations/'
txtsavepath = fdir
total_xml = os.listdir(xmlfilepath)
random.shuffle(total_xml)

num = len(total_xml)
num_list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(num_list, tv)
trainval.sort(key=int)
train = random.sample(trainval, tr)
train.sort(key=int)

val = list(set(trainval) - set(train))
test = list(set(num_list) - set(trainval))

ftrainval = open(fdir + 'trainval.txt', 'w')
ftest = open(fdir + 'test.txt', 'w')
ftrain = open(fdir + 'train.txt', 'w')
fval = open(fdir + 'val.txt', 'w')

start = cv2.getTickCount()
for i in trainval:
    name = total_xml[i][:-4] + '\n'
    ftrainval.write(name)

for i in train:
    name = total_xml[i][:-4] + '\n'
    ftrain.write(name)

for i in val:
    name = total_xml[i][:-4] + '\n'
    fval.write(name)

for i in test:
    name = total_xml[i][:-4] + '\n'
    ftest.write(name)

end = cv2.getTickCount()
during = (end - start) / cv2.getTickFrequency()
print("time: {}".format(during))

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

3.复制到images目录下,标注txt复制到labels目录下

import os
import shutil
from tqdm import tqdm

SPLIT_PATH = "F:/Project_code/yolov7-main/VOCdevkit/VOC2007/ImageSets/Main"
IMGS_PATH = "F:/Project_code/yolov7-main/VOCdevkit/VOC2007/JPEGImages"
TXTS_PATH = "F:/Project_code/yolov7-main/VOCdevkit/VOC2007/labels_copy"

TO_IMGS_PATH = 'F:/Project_code/yolov7-main/VOCdevkit/VOC2007/images'
TO_TXTS_PATH = 'F:/Project_code/yolov7-main/VOCdevkit/VOC2007/labels'

data_split = ['trainval.txt', 'test.txt']
to_split = ['train', 'val']

for index, split in enumerate(data_split):
    split_path = os.path.join(SPLIT_PATH, split)

    to_imgs_path = os.path.join(TO_IMGS_PATH, to_split[index])
    if not os.path.exists(to_imgs_path):
        os.makedirs(to_imgs_path)

    to_txts_path = os.path.join(TO_TXTS_PATH, to_split[index])
    if not os.path.exists(to_txts_path):
        os.makedirs(to_txts_path)

    f = open(split_path, 'r')
    count = 1

    for line in tqdm(f.readlines(), desc="{} is copying".format(to_split[index])):
        # 复制图片
        src_img_path = os.path.join(IMGS_PATH, line.strip() + '.jpg')
        dst_img_path = os.path.join(to_imgs_path, line.strip() + '.jpg')
        if os.path.exists(src_img_path):
            shutil.copyfile(src_img_path, dst_img_path)
        else:
            print("error file: {}".format(src_img_path))

        # 复制txt标注文件
        src_txt_path = os.path.join(TXTS_PATH, line.strip() + '.txt')
        dst_txt_path = os.path.join(to_txts_path, line.strip() + '.txt')
        if os.path.exists(src_txt_path):
            shutil.copyfile(src_txt_path, dst_txt_path)
        else:
            print("error file: {}".format(src_txt_path))

你可能感兴趣的:(Python,python,java,人工智能)