defaddImgItem(file_name, size): global image_id, id_num if file_name isNone: raise Exception('Could not find filename tag in xml file.') if size['width'] isNone: raise Exception('Could not find width tag in xml file.') if size['height'] isNone: raise Exception('Could not find height tag in xml file.')
tree = ET.parse(xml_file) root = tree.getroot() if root.tag != 'annotation': raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))
# elem is <folder>, <filename>, <size>, <object> for elem in root: current_parent = elem.tag current_sub = None object_name = None
if elem.tag == 'folder': continue
if elem.tag == 'filename': file_name = elem.text if file_name in category_set: raise Exception('file_name duplicated')
# add img item only after parse <size> tag elif current_image_id isNoneand file_name isnotNoneand size['width'] isnotNone: if file_name notin image_set: current_image_id = addImgItem(file_name, size) print('add image with {} and {}'.format(file_name, size)) else: raise Exception('duplicated image: {}'.format(file_name)) # subelem is <width>, <height>, <depth>, <name>, <bndbox> for subelem in elem: bndbox['xmin'] = None bndbox['xmax'] = None bndbox['ymin'] = None bndbox['ymax'] = None
elif current_parent == 'size': if size[subelem.tag] isnotNone: raise Exception('xml structure broken at size tag.') size[subelem.tag] = int(subelem.text)
# option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox> for option in subelem: if current_sub == 'bndbox': if bndbox[option.tag] isnotNone: raise Exception('xml structure corrupted at bndbox tag.') bndbox[option.tag] = int(option.text)
# only after parse the <object> tag if bndbox['xmin'] isnotNone: if object_name isNone: raise Exception('xml structure broken at bndbox tag') if current_image_id isNone: raise Exception('xml structure broken at bndbox tag') if current_category_id isNone: raise Exception('xml structure broken at bndbox tag') bbox = [] # x bbox.append(bndbox['xmin']) # y bbox.append(bndbox['ymin']) # w bbox.append(bndbox['xmax'] - bndbox['xmin']) # h bbox.append(bndbox['ymax'] - bndbox['ymin']) print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id, bbox)) addAnnoItem(object_name, current_image_id, current_category_id, bbox)
# 获取训练集和测试集图像的 ID 列表 train_image_ids = [img['id'] for img in train_images] test_image_ids = [img['id'] for img in test_images]
# 根据图像 ID 过滤标注信息，得到训练集和测试集的标注 train_annotations = [ann for ann in annotations if ann['image_id'] in train_image_ids] test_annotations = [ann for ann in annotations if ann['image_id'] in test_image_ids]
wechat
Alipay
