python批量修改标记完的xml文件

目标

已经打完标签的数据集，有以下需求：

• 图片太大了不方便训练，需要把尺寸调低
• 数据集移动了位置，xml中的路径需要改
   

方法

• 涉及到四个文件夹，分别是：xml原位置、xml新位置、img原位置、img新位置；
• 先批量读取img原位置所有图片，缩放后保存到img新位置；
• 批量读取xml原位置所有文件，复制一份到xml新位置，并在新位置的文件中写入更新
   

类型一：统一缩放比例

import os
import cv2
import shutil
import xml.etree.ElementTree as ET

xml_src_dir = ""    #xml原位置
xml_dst_dir = ""    #修改后xml的保存位置
img_src_dir = ""    #图片原位置
img_dst_dir = ""    #修改后图片的保存位置

img_files = os.listdir(img_src_dir)
xml_files = os.listdir(xml_src_dir)

ratio = 0.1   #图片缩放比例

# 图片缩放并保存到新路径
for img_src_file in img_files:
    img_src = cv2.imread(img_src_dir + img_src_file)
    img_dst = cv2.resize(img_src,None,fx = ratio, fy = ratio)
    cv2.imwrite(img_dst_dir + img_src_file,img_dst)

for xml_src_file in xml_files:
    if xml_src_file.endswith('.xml'):

        # 从xml原路径复制一份到新位置
        shutil.copy(xml_src_dir + xml_src_file, xml_dst_dir)

        # 修改xml中的图片尺寸
        xml_dst = ET.parse(xml_dst_dir + xml_src_file)
        w_src=float(xml_dst.find('size').find('width').text)
        h_src=float(xml_dst.find('size').find('height').text)
        xml_dst.find('size').find('width').text = str(round(w_src*ratio,1))
        xml_dst.find('size').find('height').text = str(round(h_src*ratio,1))

        # 修改xml中所有标签尺寸
        for obj in xml_dst.iter('object'):
            xmin=float(obj.find('bndbox').find('xmin').text)
            ymin=float(obj.find('bndbox').find('ymin').text)
            xmax=float(obj.find('bndbox').find('xmax').text)
            ymax=float(obj.find('bndbox').find('ymax').text)
            obj.find('bndbox').find('xmin').text=str(round(xmin*ratio,1))
            obj.find('bndbox').find('ymin').text=str(round(ymin*ratio,1))
            obj.find('bndbox').find('xmax').text=str(round(xmax*ratio,1))
            obj.find('bndbox').find('ymax').text=str(round(ymax*ratio,1))
        xml_dst.write(xml_dst_dir + xml_src_file)

        # 修改xml中的路径
        tree = ET.ElementTree(file = xml_dst_dir + xml_src_file)
        root = tree.getroot()
        root[0].text = 'image'
        part = os.path.splitext(xml_src_file)[0]
        part1 = part + '.jpg'
        root[2].text = img_dst_dir + part1
        tree.write(xml_dst_dir + xml_src_file)

 

类型二：统一尺寸

import os
import cv2
import shutil
import xml.etree.ElementTree as ET

xml_src_dir = ""    #xml原位置
xml_dst_dir = ""    #修改后xml的保存位置
img_src_dir = ""    #图片原位置
img_dst_dir = ""    #修改后图片的保存位置

img_files = os.listdir(img_src_dir)
xml_files = os.listdir(xml_src_dir)

# 统一尺寸
w_dst = 512
h_dst = 288

# 图片缩放并保存到新路径
for img_src_file in img_files:
    img_src = cv2.imread(img_src_dir + img_src_file)
    ratio_x = w_dst/img_src.shape[1]
    ratio_y = h_dst/img_src.shape[0]
    img_dst = cv2.resize(img_src,None,fx = ratio_x, fy = ratio_y)
    cv2.imwrite(img_dst_dir + img_src_file,img_dst)

for xml_src_file in xml_files:
    if xml_src_file.endswith('.xml'):

        # 从xml原路径复制一份到新位置
        shutil.copy(xml_src_dir + xml_src_file, xml_dst_dir)

        # 修改xml中的图片尺寸
        xml_dst = ET.parse(xml_dst_dir + xml_src_file)
        w_src=float(xml_dst.find('size').find('width').text)
        h_src=float(xml_dst.find('size').find('height').text)
        xml_dst.find('size').find('width').text = str(w_dst)
        xml_dst.find('size').find('height').text = str(h_dst)
        
        ratio_x = w_dst/w_src
        ratio_y = h_dst/h_src
        
        # 修改xml中所有标签尺寸
        for obj in xml_dst.iter('object'):
            xmin=float(obj.find('bndbox').find('xmin').text)
            ymin=float(obj.find('bndbox').find('ymin').text)
            xmax=float(obj.find('bndbox').find('xmax').text)
            ymax=float(obj.find('bndbox').find('ymax').text)
            obj.find('bndbox').find('xmin').text=str(xmin*ratio_x)
            obj.find('bndbox').find('ymin').text=str(ymin*ratio_y)
            obj.find('bndbox').find('xmax').text=str(xmax*ratio_x)
            obj.find('bndbox').find('ymax').text=str(ymax*ratio_y)
        xml_dst.write(xml_dst_dir + xml_src_file)

        # 修改xml中的路径
        tree = ET.ElementTree(file = xml_dst_dir + xml_src_file)
        root = tree.getroot()
        root[0].text = 'image'
        part = os.path.splitext(xml_src_file)[0]
        part1 = part + '.jpg'
        root[2].text = img_dst_dir + part1
        tree.write(xml_dst_dir + xml_src_file)