Python读取xml文件后，裁剪标注图片，并扩容数据

0、引子

因为想用yolov5算法训练自己数据集识别数字“0-9”，一开始用labelImg标注了图片，生成了大量的xml文件。因为图片中0，1比较多，而其他数字偏少，标注到后面，就忽略了大量的0，1。后面发现，漏标注会导致训练时把目标识别成背景，严重影响算法识别的准确性。然后，我也不想重新去标注图片了，就想着写个Python程序根据xml文件，按照标注框，把目标都裁剪出来。

1、裁剪图片

首先是根据xml文件把对应标注图片，按标注框，裁剪出来。我在参考博客的基础上实现了裁剪图片按类别保存到对应文件夹里面，并在该类别下按顺序编号

# 导入模块
import cv2
import xml.etree.ElementTree as ET
import os
from pathlib import Path
import numpy as np
import random

# 原图片、标签文件、裁剪图片路径
img_path = r'D:\yolov5-3.1\cut\c_1'
xml_path = r'D:\yolov5-3.1\cut\xml'
obj_img_path = r'D:\yolov5-3.1\cut\c_3'

# 声明一个空字典用于储存裁剪图片的类别及其数量
Numpic = {}

# 把原图片裁剪后，按类别新建文件夹保存，并在该类别下按顺序编号
for img_file in os.listdir(img_path):
    if img_file[-4:] in ['.png', '.jpg']:  # 判断文件是否为图片格式
        img_filename = os.path.join(img_path, img_file)  # 将图片路径与图片名进行拼接
        img_cv = cv2.imread(img_filename)  # 读取图片

        img_name = (os.path.splitext(img_file)[0])  # 分割出图片名，如“000.png” 图片名为“000”
        xml_name = xml_path + '\\' + '%s.xml' % img_name  # 利用标签路径、图片名、xml后缀拼接出完整的标签路径名

        if os.path.exists(xml_name):  # 判断与图片同名的标签是否存在，因为图片不一定每张都打标
            root = ET.parse(xml_name).getroot()  # 利用ET读取xml文件
            for obj in root.iter('object'):  # 遍历所有目标框
                name = obj.find('name').text  # 获取目标框名称，即label名
                xmlbox = obj.find('bndbox')  # 找到框目标
                x0 = xmlbox.find('xmin').text  # 将框目标的四个顶点坐标取出
                y0 = xmlbox.find('ymin').text
                x1 = xmlbox.find('xmax').text
                y1 = xmlbox.find('ymax').text

                obj_img = img_cv[int(y0):int(y1), int(x0):int(x1)]  # cv2裁剪出目标框中的图片

                Numpic.setdefault(name, 0)  # 判断字典中有无当前name对应的类别，无则新建
                Numpic[name] += 1  # 当前类别对应数量 + 1
                my_file = Path(obj_img_path + '\\' + name)  # 判断当前name对应的类别有无文件夹
                if 1 - my_file.is_dir():  # 无则新建
                    os.mkdir(obj_img_path + '\\' + str(name))

                cv2.imwrite(obj_img_path + '\\' + name + '\\' + '%04d' % (Numpic[name]) + '.jpg',
                            obj_img)  # 保存裁剪图片，图片命名4位，不足补0

2、图片扩容

只是把标注框裁剪出来，还会有一个问题就是，每个类别的数量不一致，0，1的图片多，其他数字少，作为训练集可能不太好。我想，要是每个类别的图片数量都一致就好了。于是我继续把裁剪图片进行扩容，这里只是通过给图片增加噪点来扩容。

# 新建一个图片加噪点的函数
def random_noise(image,noise_num):
    img_noiseimg = cv2.imread(image) # 读取图片
    rows, cols, chn = img_noise.shape
 
    for i in range(noise_num):
        x = np.random.randint(0, rows)#随机生成指定范围的整数
        y = np.random.randint(0, cols)
        img_noise[x, y, :] = 0 # 0代表黑色，255代表白色
    return img_noise

# 图片扩容
max_Numpic = max(Numpic.values()) # 提取裁剪图片中，类别下数量最大值
for name in Numpic:# 遍历每一个类别
    for i in range (Numpic[name] + 1, max_Numpic + 1):# 把其余类别的图片数量扩充到，与数量值最大的类别相等（我的数据集里面“0”这个类别数量是最多的）
        Noisenum = random.randint(1, 20)# 生成随机的噪点数
        Num = random.randint(1, Numpic[name])# 随机选择该类别下已存在的一个图片
        Noicepic = random_noise(obj_img_path + '\\' + name + '\\' + '%04d' % Num + '.jpg', Noisenum)# 给图片加噪点
        cv2.imwrite(obj_img_path + '\\' + name + '\\' + '%04d' % (i) + '.jpg', Noicepic)# 保存图片

3、调用训练集

接下来就是怎样调用这些裁剪图片作为yolov5算法的训练数据集，这里先挖个坑，下周解决了再继续写。

参考：python利用文件夹下xml格式标签文件批量裁剪出图片中的目标（文件夹、图片名称、目标框数量无限制，逐行注释）