VOC数据集转化为COCO数据集
代码如下:
#!/usr/bin/python
# pip install lxml
import os
import json
import xml.etree.ElementTree as ET
from tqdm import tqdm
START_BOUNDING_BOX_ID = 1
PRE_DEFINE_CATEGORIES = {} # 类别id对应键值, 如果有类别关, 需要填写此项, 即VOC中names对应的类别id如 1, 2, 3...
IMAGE_ID = 0 # 全局图片ID
# If necessary, pre-define category and its id
# PRE_DEFINE_CATEGORIES = {"aeroplane": 1, "bicycle": 2, "bird": 3, "boat": 4,
# "bottle":5, "bus": 6, "car": 7, "cat": 8, "chair": 9,
# "cow": 10, "diningtable": 11, "dog": 12, "horse": 13,
# "motorbike": 14, "person": 15, "pottedplant": 16,
# "sheep": 17, "sofa": 18, "train": 19, "tvmonitor": 20}
def create_xml_list(xml_dir, output_path):
xml_list = os.listdir(xml_dir)
with open(output_path, 'w') as f:
for i in xml_list:
if i[-3:] == 'xml':
f.write(str(i) + '\n')
def get(root, name):
vars = root.findall(name)
return vars
def get_and_check(root, name, length):
vars = root.findall(name)
if len(vars) == 0:
raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))
if length > 0 and len(vars) != length:
raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
if length == 1:
vars = vars[0]
return vars
def get_filename_as_int(filename):
try:
filename = os.path.splitext(filename)[0]
global IMAGE_ID
IMAGE_ID = IMAGE_ID + 1
return IMAGE_ID
# return int(filename)
except:
raise NotImplementedError('Filename %s is supposed to be an integer.' % (filename))
def convert(xml_list, xml_dir, json_file):
list_fp = open(xml_list, 'r')
json_dict = {"images": [], "type": "instances", "annotations": [],
"categories": []}
categories = PRE_DEFINE_CATEGORIES
bnd_id = START_BOUNDING_BOX_ID
for line in tqdm(list_fp):
line = line.strip()
# print("Processing %s" % (line))
xml_f = os.path.join(xml_dir, line)
tree = ET.parse(xml_f)
root = tree.getroot()
path = get(root, 'path')
if len(path) == 1:
filename = os.path.basename(path[0].text)
elif len(path) == 0:
filename = get_and_check(root, 'filename', 1).text
else:
raise NotImplementedError('%d paths found in %s' % (len(path), line))
## The filename must be a number
image_id = get_filename_as_int(filename)
size = get_and_check(root, 'size', 1)
width = int(get_and_check(size, 'width', 1).text)
height = int(get_and_check(size, 'height', 1).text)
image = {'file_name': filename, 'height': height, 'width': width,
'id': image_id}
json_dict['images'].append(image)
## Cruuently we do not support segmentation
# segmented = get_and_check(root, 'segmented', 1).text
# assert segmented == '0'
for obj in get(root, 'object'):
category = get_and_check(obj, 'name', 1).text
# if category not in categories:
# new_id = len(categories)
# categories[category] = new_id
# category_id = categories[category]
bndbox = get_and_check(obj, 'bndbox', 1)
xmin = float(get_and_check(bndbox, 'xmin', 1).text) # 使用 float 保留原精度
ymin = float(get_and_check(bndbox, 'ymin', 1).text)
xmax = float(get_and_check(bndbox, 'xmax', 1).text)
ymax = float(get_and_check(bndbox, 'ymax', 1).text)
assert (xmax > xmin)
assert (ymax > ymin)
o_width = float(xmax - xmin)
o_height = float(ymax - ymin)
ann = {'area': o_width * o_height, 'iscrowd': 0, 'image_id':
image_id, 'bbox': [xmin, ymin, o_width, o_height],
'category_id': int(category), 'id': bnd_id, 'ignore': 0,
'segmentation': []}
json_dict['annotations'].append(ann)
bnd_id = bnd_id + 1
for cate, cid in categories.items():
cat = {'supercategory': 'none', 'id': cid, 'name': cate}
json_dict['categories'].append(cat)
json_fp = open(json_file, 'w')
json_str = json.dumps(json_dict)
json_fp.write(json_str)
json_fp.close()
list_fp.close()
if __name__ == '__main__':
# if len(sys.argv) <= 1:
# print('3 auguments are need.')
# print('Usage: %s XML_LIST.txt XML_DIR OUTPU_JSON.json'%(sys.argv[0]))
# exit(1)
#
# convert(sys.argv[1], sys.argv[2], sys.argv[3])
xml_path = 'D:/Users/Desktop/GradProject/datasets/test_xml2'
xml_list = 'D:/Users/Desktop/xml_list.txt'
json_path = 'D:/Users/Desktop/test_annotations.json'
create_xml_list(xml_path, xml_list)
convert(xml_list, xml_path, json_path)
转换流程:
- 先将要转化的所有.xml文件名保存在xml_list.txt列表中。
- 遍历 xml_list.txt 的文件进行转化