PASCAL VOC标注文件解析代码

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/hzhj2007/article/details/80825365

功能函数：从检测的xml文件中提取物体，并保存。

• 看下xml的文件内容：

<?xml version="1.0" ?>
<root>
	<filename>2464-1</filename>
	<source>2sc</source>
	<size>
		<width>550</width>
		<height>375</height>
		<depth>3</depth>
	</size>
	<object>
		<name>yinzhangbbox</name>
		<bndbox>
			<xmin>21</xmin>
			<ymin>235</ymin>
			<xmax>144</xmax>
			<ymax>358</ymax>
		</bndbox>
	</object>
	<object>
		<name>v_title</name>
		<bndbox>
			<xmin>103</xmin>
			<ymin>17</ymin>
			<xmax>440</xmax>
			<ymax>55</ymax>
		</bndbox>
	</object>
</root>

• 功能函数

#encoding:utf-8
import os,sys,cv2
import copy
import argparse
import xml.etree.ElementTree as ET

def get_file(path):
    allfile=[]
    for dirpath,dirnames,filenames in os.walk(path):
        for dir in dirnames:
            allfile.append(os.path.join(dirpath,dir))
        for name in filenames:
            if name.find('.xml') == -1:
                continue
            allfile.append(os.path.join(dirpath, name))
    return allfile

def crop_img(img, boxes):
    img_crops = []
    for box in boxes:
        x1, y1, x2, y2 = box
        img_crop = img[y1:y2, x1:x2,:]
        img_crops.append(img_crop)
    return img_crops

def get_img_path(xml_path):
    xml_path_cp = xml_path.replace('Annotations', 'JPEGImages')
    return xml_path_cp.replace('.xml', '.jpg')

def save_all_crops(img_crops, file_name, save_dir):
    base_name = os.path.basename(file_name)
    for idx, img_crop in enumerate(img_crops):
        rst_name = os.path.join(save_dir, base_name[:base_name.find('.jpg')] + '_' + str(idx) + base_name[base_name.find('.jpg'):])
        cv2.imwrite(rst_name, img_crop)

def get_all_crop(xml_path):
    rtn_boxes = []
    et = ET.parse(xml_path)
    element = et.getroot()

    element_objs = element.findall('object')
    element_filename = element.find('filename').text
    element_width = int(element.find('size').find('width').text)
    element_height = int(element.find('size').find('height').text)
    for obj in element_objs:
        x1 = obj.find('bndbox').find('xmin').text
        y1 = obj.find('bndbox').find('ymin').text
        x2 = obj.find('bndbox').find('xmax').text
        y2 = obj.find('bndbox').find('ymax').text
        rtn_boxes.append([x1, y1, x2, y2])
    return rtn_boxes

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--xml_path', default='./data/')
    parser.add_argument('--save_dir', default='./my_save/')

    return parser.parse_args()

if __name__ == '__main__':
    args = get_args()
    xmls_path = get_file(args.xml_path)
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    for xml_path in xmls_path:
        boxes = get_all_crop(xml_path)
        img_path = get_img_path(xml_path)
        img = cv2.imread(img_path)
        img_crops = crop_img(img, boxes)

        save_all_crops(img_crops, img_path, args.save_dir)

• 运行脚本命令：

python  parse_xml_to_save.py --xml_path='PATH_TO/Annotations/'  \
                             --save_dir='./crop_from_xml'