用Python读取xml文件后,裁剪标注图片和扩容数据

本文的文字及图片来源于网络,仅供学习、交流使用,不具有任何商业用途,,版权归原作者所有,如有问题请及时联系我们以作处理

作者:Gendan

原文https://segmentfault.com/a/1190000038690313

想要获取更多Python学习资料,了解更多关于Python的知识,可以加Q群630390733踊跃发言,大家一起来学习讨论吧!

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

1、裁剪图片

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

导入模块import cv2import xml.etree.ElementTree as ETimport osfrom pathlib import Pathimport numpy as npimport random

原图片、标签文件、裁剪图片路径

img_path = r'D:yolov5-3.1cutc_1'

xml_path = r'D:yolov5-3.1cutxml'

obj_img_path = r'D:yolov5-3.1cutc_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、图片扩容

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

新建一个图片加噪点的函数

def random_noise(image,noise_num):

img_noiseimg = cv2.imread(image) # 读取图片

ows, 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)# 保存图片

你可能感兴趣的:(用Python读取xml文件后,裁剪标注图片和扩容数据)