【YOLOV5学习记录,txt文件直接划分比例】


平时学习记录一些日常,非原创。

一、图片标注

使用开源工具 labelImg 对图片进行标注,导出的数据集格式为PASCAL VOC,待数据标注完成后,可以看到文件夹是下面这样,标注文件 txt和图片文件混在了一起

【YOLOV5学习记录,txt文件直接划分比例】_第1张图片

二、YOLO数据格式

大家都知道,yolov5训练时使用的数据集结构是这样的:

├─test
│  ├─images
│  └─labels
├─train
│  ├─images
│  └─labels
└─valid
    ├─images
    └─labels

因此,我们可以将图片文件和对应的txt标签文件进行一次划分,首先创建外层的trainvalidtest文件夹,然后在每个文件夹底下都分别创建imageslabels文件夹。

注意:需要创建相应的文件夹,运行脚本之前代码目录结构如下。split.py和外层属于同一目录,同级目录


├─images
│  │ 
│  └─.jpg 图片
├─labels
│  │ 
│  └─.txt 文件
├─test
│  ├─images
│  └─labels
├─train
│  ├─images
│  └─labels
├─valid
│  ├─images
│  └─labels
│  
└─split.py


接下来,可以使用下面的脚本,将图片和标签文件按照比例进行划分:

import os
import shutil
import random

# 训练集、验证集和测试集的比例分配
test_percent = 0.1
valid_percent = 0.2
train_percent = 0.7

# 标注文件的路径
image_path = 'data_t/images'
label_path = 'data_t/labels'

images_files_list = os.listdir(image_path)
labels_files_list = os.listdir(label_path)
print('images files: {}'.format(images_files_list))
print('labels files: {}'.format(labels_files_list))
total_num = len(images_files_list)
print('total_num: {}'.format(total_num))

test_num = int(total_num * test_percent)
valid_num = int(total_num * valid_percent)
train_num = int(total_num * train_percent)

# 对应文件的索引
test_image_index = random.sample(range(total_num), test_num)
valid_image_index = random.sample(range(total_num), valid_num)
train_image_index = random.sample(range(total_num), train_num)

for i in range(total_num):
    print('src image: {}, i={}'.format(images_files_list[i], i))
    # 将图片和标签文件拷贝到对应文件夹下
    if i in test_image_index:
        # Test
        shutil.copyfile('data_t\\images\\{}'.format(images_files_list[i]),
                        'data_t\\test\\images\\{}'.format(images_files_list[i]))
        shutil.copyfile('data_t\\labels\\{}'.format(labels_files_list[i]),
                        'data_t\\test\\labels\\{}'.format(labels_files_list[i]))
    elif i in valid_image_index:
        # Valid
        shutil.copyfile('data_t//images//{}'.format(images_files_list[i]),
                        'data_t//valid//images//{}'.format(images_files_list[i]))
        shutil.copyfile('data_t//labels//{}'.format(labels_files_list[i]),
                        'data_t//valid//labels//{}'.format(labels_files_list[i]))
    else:
        # Train
        shutil.copyfile('data_t/images/{}'.format(images_files_list[i]),
                        'data_t/train/images/{}'.format(images_files_list[i]))
        shutil.copyfile('data_t/labels/{}'.format(labels_files_list[i]),
                        'data_t/train/labels/{}'.format(labels_files_list[i]))
# 也可以用绝对路径
# 文件目录路径,三种斜杠方式均可。但是最好用一致的。“/”和"\\"优先。自己记录这种方式而已

 执行代码后,可以看到类似文件层级结构:


├─split.py
├─images
│  │ 
│  └─.jpg 图片
├─labels
│  │ 
│  └─.txt 文件
├─test
│  ├─images
│  │      1234565343231.jpg
│  │      1559035146628.jpg
│  │      2019032210151.jpg
│  │
│  └─labels
│          1234565343231.txt
│          1559035146628.txt
│          2019032210151.txt
│
├─train
│  ├─images
│  │      1213211.jpg
│  │      12i4u33112.jpg
│  │      1559092537114.jpg
│  │
│  └─labels
│          1213211.txt
│          12i4u33112.txt
│          1559092537114.txt
│
└─valid
    ├─images
    │      120131247621.jpg
    │      124iuy311.jpg
    │      1559093141383.jpg
    │
    └─labels
            120131247621.txt
            124iuy311.txt
            1559093141383.txt

至此,数据集就真正准备好了。 

2.YOLO转VOC

如果拿到了txt的标注,但是需要使用VOC,也需要进行转换。看下面这个脚本,注释写在代码中

import os
import xml.etree.ElementTree as ET
from PIL import Image
import numpy as np

# 图片文件夹,后面的/不能省
img_path = 'images/'

# txt文件夹,后面的/不能省
labels_path = 'labels/'

# xml存放的文件夹,后面的/不能省
annotations_path = 'Annotations/'

labels = os.listdir(labels_path)

# 类别 这里必须修改
classes = ["hat"]

# 图片的高度、宽度、深度
sh = sw = sd = 0

def write_xml(imgname, sw, sh, sd, filepath, labeldicts):
    '''
    imgname: 没有扩展名的图片名称
    '''

    # 创建Annotation根节点
    root = ET.Element('Annotation')

    # 创建filename子节点,无扩展名                 
    ET.SubElement(root, 'filename').text = str(imgname)        

    # 创建size子节点 
    sizes = ET.SubElement(root,'size')                                      
    ET.SubElement(sizes, 'width').text = str(sw)
    ET.SubElement(sizes, 'height').text = str(sh)
    ET.SubElement(sizes, 'depth').text = str(sd) 

    for labeldict in labeldicts:
        objects = ET.SubElement(root, 'object')                 
        ET.SubElement(objects, 'name').text = labeldict['name']
        ET.SubElement(objects, 'pose').text = 'Unspecified'
        ET.SubElement(objects, 'truncated').text = '0'
        ET.SubElement(objects, 'difficult').text = '0'
        bndbox = ET.SubElement(objects,'bndbox')
        ET.SubElement(bndbox, 'xmin').text = str(int(labeldict['xmin']))
        ET.SubElement(bndbox, 'ymin').text = str(int(labeldict['ymin']))
        ET.SubElement(bndbox, 'xmax').text = str(int(labeldict['xmax']))
        ET.SubElement(bndbox, 'ymax').text = str(int(labeldict['ymax']))
    tree = ET.ElementTree(root)
    tree.write(filepath, encoding='utf-8')


for label in labels:
    with open(labels_path + label, 'r') as f:
        img_id = os.path.splitext(label)[0]
        contents = f.readlines()
        labeldicts = []
        for content in contents:
            # 这里要看你的图片格式了,我这里是jpg,注意修改
            img = np.array(Image.open(img_path + label.strip('.txt') + '.jpg'))

            # 图片的高度和宽度
            sh, sw, sd = img.shape[0], img.shape[1], img.shape[2]
            content = content.strip('\n').split()
            x = float(content[1])*sw
            y = float(content[2])*sh
            w = float(content[3])*sw
            h = float(content[4])*sh

            # 坐标的转换,x_center y_center width height -> xmin ymin xmax ymax
            new_dict = {'name': classes[int(content[0])],
                        'difficult': '0',
                        'xmin': x+1-w/2,                     
                        'ymin': y+1-h/2,
                        'xmax': x+1+w/2,
                        'ymax': y+1+h/2
                        }
            labeldicts.append(new_dict)
        write_xml(img_id, sw, sh, sd, annotations_path + label.strip('.txt') + '.xml', labeldicts)

执行上述脚本,就可以在Annotations看到转换后的xml文件了。

3.VOC数据集格式

首先,按照VOC2007的数据集格式要求,分别创建文件夹VOCdevkitVOC2007AnnotationsImageSetsMainJPEGImages,它们的层级结构如下所示:

└─VOCdevkit
    └─VOC2007
        ├─Annotations
        ├─ImageSets
        │  └─Main
        └─JPEGImages

其中,Annotations用来存放xml标注文件,JPEGImages用来存放图片文件,而ImageSets/Main存放几个txt文本文件,文件的内容是训练集、验证集和测试集中图片的名称。Main最开始为空目录,文本文件之后才生成。

接下来新建一个脚本,把它放在VOCdevkit/VOC2007文件夹下,命名为test.py,脚本的内容如下:

import os
import random

# 训练集和验证集的比例分配
trainval_percent = 0.1
train_percent = 0.9

# 标注文件的路径
xmlfilepath = 'Annotations'

# 生成的txt文件存放路径
txtsavepath = 'ImageSets\Main'
total_xml = os.listdir(xmlfilepath)

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

ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)

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

执行这个脚本,结束后,在ImageSets/Main下生成了4个txt文件。

test.txt                train.txt                 trainval.txt                 val.txt

这4个文件的格式都是一样的,文件的内容是对应图片名称去掉扩展名(与xml标注文件去掉.xml一致)的结果 

import xml.etree.ElementTree as ET
import os

sets = ['train', 'val', 'test','trainval']  # 如果你的Main文件夹没有test.txt,就删掉'test'
# classes = ["a", "b"]   # 改成自己的类别,VOC数据集有以下20类别
classes = ["Screw", 'Border']  # class names
# 替换博客中的反斜杠
# abs_path = os.getcwd()
abs_path = os.getcwd().replace('\\', '/')
print(abs_path)


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    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 convert_annotation(image_id):
    in_file = open(abs_path + '/VOCdevkit/VOC2007/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open(abs_path + '/VOCdevkit/VOC2007/labels/%s.txt' % (image_id), '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
        # difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        # if cls not in classes or int(difficult) == 1:
        if cls not in classes:
            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))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


for image_set in sets:
    if not os.path.exists(abs_path + '/VOCdevkit/VOC2007/labels/'):
        os.makedirs(abs_path + '/VOCdevkit/VOC2007/labels/')
    # 再博主源码上 空格换成换行符。不然文件名带有空格的他分成了2个
    # https://blog.csdn.net/retainenergy/article/details/124613553?spm=1001.2101.3001.6650.5&utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7ERate-5-124613553-blog-120022845.pc_relevant_aa&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7ERate-5-124613553-blog-120022845.pc_relevant_aa
    image_ids = open(abs_path + '/VOCdevkit/VOC2007/ImageSets/Main/%s.txt' % (image_set)).read().strip().split('\n')
    list_file = open(abs_path + '/VOCdevkit/VOC2007/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        # 如果你的数据集一开始存放图片的文件夹就是images的话,此处就不需要改动,
        # 但是在这要把最下面的读取图片的文件夹路径对应成images文件夹
        list_file.write(abs_path + '/VOCdevkit/VOC2007/JPEGImages/%s.jpg\n' % (image_id))  # 要么自己补全路径,只写一半可能会报错
        convert_annotation(image_id)
    list_file.close()

链接地址: VOC与YOLO数据格式的相互转换 - 知乎环境windows 10 64位python 3.7labelImg图片标注这里使用开源工具 labelImg 对图片进行标注,导出的数据集格式为PASCAL VOC,待数据标注完成后,可以看到文件夹是下面这个样子的,标注文件xml和图片文件混在了一起…https://zhuanlan.zhihu.com/p/383660741

VOC与YOLO数据格式的相互转换-一一网络


你可能感兴趣的:(python,图像处理)