58 lines
2.1 KiB
Python
Executable File
58 lines
2.1 KiB
Python
Executable File
"""
|
||
author:Hatcher
|
||
email:hcwang@semi.ac.cn
|
||
use:将aitod小目标检测数据集按照类别分开
|
||
input：
|
||
voc
|
||
train_images
|
||
train_annos
|
||
....
|
||
test&val
|
||
output:
|
||
voc_class
|
||
classname1
|
||
images
|
||
annos
|
||
classname2
|
||
iamges
|
||
annos
|
||
.....
|
||
|
||
"""
|
||
import os
|
||
from tqdm import tqdm
|
||
import xml.dom.minidom
|
||
import shutil
|
||
source_path_root = "/data4/Hatcher/smallobject/voc/test_annos/"
|
||
det_path_root = "/data4/Hatcher/smallobject/voc_class/"
|
||
class_names = ["airplane"]#os.listdir(det_path_root)
|
||
for class_name in class_names:
|
||
for source_path_anno in tqdm(os.listdir(source_path_root)):
|
||
source_path = os.path.join(source_path_root,source_path_anno)
|
||
dom = xml.dom.minidom.parse(source_path)
|
||
root = dom.documentElement
|
||
objects=dom.getElementsByTagName("object")
|
||
for object in objects:
|
||
bndbox = object.getElementsByTagName('bndbox')[0]
|
||
xmin = bndbox.getElementsByTagName('xmin')[0]
|
||
ymin = bndbox.getElementsByTagName('ymin')[0]
|
||
xmax = bndbox.getElementsByTagName('xmax')[0]
|
||
ymax = bndbox.getElementsByTagName('ymax')[0]
|
||
#print(xmin.childNodes[0].data)
|
||
#print(type(xmin.childNodes[0].data))
|
||
xmin_data=int(xmin.childNodes[0].data)
|
||
ymin_data=int(ymin.childNodes[0].data)
|
||
xmax_data=int(xmax.childNodes[0].data)
|
||
ymax_data=int(ymax.childNodes[0].data)
|
||
label_name=object.getElementsByTagName('name')[0].childNodes[0].data
|
||
if not (label_name==class_name):
|
||
dom.documentElement.removeChild(object)
|
||
objects=dom.getElementsByTagName("object")
|
||
if not (len(objects)==0):
|
||
#存储xml，复制图像
|
||
xml_path = os.path.join(det_path_root,class_name,"annos",source_path_anno)
|
||
fp = open(xml_path, 'w', encoding='utf-8')
|
||
dom.writexml(fp, indent='', addindent='\t', newl='\n', encoding='utf-8')
|
||
fp.close()
|
||
shutil.copy(source_path.replace("test_annos","test_images").replace(".xml",".png"),os.path.join(det_path_root,class_name,"images",source_path_anno.replace(".xml",".png")))
|