VOC格式标签转YOLO格式标签
'''
Author: hollis23
Date: 2021-10-22 16:39:24
LastEditTime: 2021-10-22 16:49:30
LastEditors: Please set LastEditors
Description: In User Settings Edit
'''
import xml. etree. ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os. path import join
import random
import shutil
'''
1. 先将xml标签文件放到Anootations文件夹下,images放到JPEGImages文件夹下
'''
work_sapce_dir = os. path. join( tt100k_parent_dir, "VOCdevkit/" )
if not os. path. isdir( work_sapce_dir) :
os. mkdir( work_sapce_dir)
work_sapce_dir = os. path. join( work_sapce_dir, "VOC2007/" )
if not os. path. isdir( work_sapce_dir) :
os. mkdir( work_sapce_dir)
jpeg_images_path = os. path. join( work_sapce_dir, 'JPEGImages' )
annotations_path = os. path. join( work_sapce_dir, 'Annotations' )
if not os. path. isdir( jpeg_images_path) :
os. mkdir( jpeg_images_path)
if not os. path. isdir( annotations_path) :
os. mkdir( annotations_path)
'''
2.根据xml文件的名字,将数据划分为训练集,验证集和测试集,并将生成的trainval.txt
val.txt, test.txt放到ImageSets/Main文件夹下(这几个txt文件中的内容都是图像/xml
标签的去后缀名称,即id)
'''
root_dir = 'I:/data/VOCdevkit/VOC2007/'
trainval_percent = 0.8
train_percent = 0.7
xmlfilepath = root_dir + 'Annotations'
txtsavepath = root_dir + '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 ( root_dir + 'ImageSets/Main/trainval.txt' , 'w' , encoding= 'utf-8' )
ftest = open ( root_dir + 'ImageSets/Main/test.txt' , 'w' , encoding= 'utf-8' )
ftrain = open ( root_dir + 'ImageSets/Main/train.txt' , 'w' , encoding= 'utf-8' )
fval = open ( root_dir + 'ImageSets/Main/val.txt' , 'w' , encoding= 'utf-8' )
for i in list :
name = total_xml[ i] [ : - 4 ] + '\n'
if i in trainval:
ftrainval. write( name)
if i in train:
ftrain. write( name)
else :
fval. write( name)
else :
ftest. write( name)
ftrainval. close( )
ftrain. close( )
fval. close( )
ftest. close( )
'''
3.在VOCDevkit目录格式下,将xml标签格式转化为yolo的.txt文件格式,需要修改类别列表
'''
sets= [ ( '2007' , 'trainval' ) , ( '2007' , 'test' ) ]
classes = [ "i2" , "i4" , "i5" , "il100" , "il60" , "il80" , "io" , "ip" , "p10" , "p11" , "p12" , "p19" , "p23" , "p26" , "p27" , "p3" , "p5" , "p6" , "pg" , "ph4" , "ph4.5" , "ph5" , "pl100" , "pl120" , "pl20" , "pl30" , "pl40" ,
"pl5" , "pl50" , "pl60" , "pl70" , "pl80" , "pm20" , "pm30" , "pm55" , "pn" , "pne" , "po" , "pr40" , "w13" , "w32" , "w55" , "w57" , "w59" , "wo" ]
def convert ( size, box) :
dw = 1 . / size[ 0 ]
dh = 1 . / 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 convert_annotation ( year, image_id) :
in_file = open ( '/data/datasets/TT100K_45/VOCdevkit/VOC%s/Annotations/%s.xml' % ( year, image_id) )
out_file = open ( '/data/datasets/TT100K_45/VOCdevkit/VOC%s/labels/%s.txt' % ( year, 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
cls = obj. find( 'name' ) . text
if cls not in classes or int ( difficult) == 1 :
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) )
bb = convert( ( w, h) , b)
out_file. write( str ( cls_id) + " " + " " . join( [ str ( a) for a in bb] ) + '\n' )
wd = '/data/datasets/TT100K_45/'
for year, image_set in sets:
if not os. path. exists( '/data/datasets/TT100K_45/VOCdevkit/VOC%s/labels/' % ( year) ) :
os. makedirs( '/data/datasets/TT100K_45/VOCdevkit/VOC%s/labels/' % ( year) )
image_ids = open ( '/data/datasets/TT100K_45/VOCdevkit/VOC%s/ImageSets/Main/%s.txt' % ( year, image_set) ) . read( ) . strip( ) . split( )
list_file = open ( '%s_%s.txt' % ( year, image_set) , 'w' )
for image_id in image_ids:
list_file. write( '%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n' % ( wd, year, image_id) )
convert_annotation( year, image_id)
list_file. close( )
'''
4. yolo格式标签划分训练集,验证机和测试集可以依据voc格式下的划分方法,即利用trainval.txt
val.txt, test.txt中的id,将相应的图片和yolo的txt标签通过shutil.copyfile放到yolo格式的
目录下
'''
root_image_path = '/data/datasets/TT100K_45/VOCdevkit/VOC2007/JPEGImages/'
root_label_path = '/data/datasets/TT100K_45/VOCdevkit/VOC2007/labels/'
image_path = '/data/datasets/TT100K_45/VOCdevkit/VOC2007/YOLO/'
label_path = '/data/datasets/TT100K_45/VOCdevkit/VOC2007/YOLO/'
images_files_list = os. listdir( root_image_path)
labels_files_list = os. listdir( root_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) )
train_file = '/data/datasets/TT100K_45/VOCdevkit/VOC2007/ImageSets/Main/trainval.txt'
test_file = '/data/datasets/TT100K_45/VOCdevkit/VOC2007/ImageSets/Main/test.txt'
with open ( train_file, 'r' ) as f:
train_id = f. readlines( )
for id in train_id:
id = id . strip( )
shutil. copyfile( root_image_path + id + '.jpg' , image_path + 'train/images/' + id + '.jpg' )
shutil. copyfile( root_label_path + id + '.txt' , label_path + 'train/labels/' + id + '.txt' )
with open ( test_file, 'r' ) as f1:
test_id = f1. readlines( )
for id in test_id:
id = id . strip( )
shutil. copyfile( root_image_path + id + '.jpg' , image_path + 'test/images/' + id + '.jpg' )
shutil. copyfile( root_label_path + id + '.txt' , label_path + 'test/labels/' + id + '.txt' )