import cv2
import xml.etree.ElementTree as ET
import os
import numpy as np
from tqdm import tqdm
from PIL import Image
# 数据扩增
# 读xml文档返回每个目标边界框信息
def read_xml_annotation(root, image_id): # (源xml文件夹地址和文件夹内所有图像名称)
in_file = open(os.path.join(root, image_id))
tree = ET.parse(in_file)
root = tree.getroot()
bndboxlist = []
for object in root.findall('object'): # 找到root节点下的所有country节点
bndbox = object.find('bndbox') # 子节点下节点rank的值
xmin = int(bndbox.find('xmin').text)
xmax = int(bndbox.find('xmax').text)
ymin = int(bndbox.find('ymin').text)
ymax = int(bndbox.find('ymax').text)
bndboxlist.append([xmin, ymin, xmax, ymax])
bndbox = root.find('object').find('bndbox')
return bndboxlist
# 更改xml信息将裁剪至每个小部分的对应目标边界框进行提取
def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
in_file = open(os.path.join(root, str(image_id) + '.xml')) # 这里root分别由两个意思
tree = ET.parse(in_file)
# 修改增强后的xml文件中的filename
elem = tree.find('filename')
elem.text = (str(id) + '.jpg')
xmlroot = tree.getroot()
# 修改增强后的xml文件中的path
elem = tree.find('path')
if elem != None:
elem.text = (saveroot + str(id) + '.jpg')
index = 0
for object in xmlroot.findall('object'): # 找到root节点下的所有country节点
bndbox = object.find('bndbox') # 子节点下节点rank的值
new_xmin = new_target[index][0]
new_ymin = new_target[index][1]
new_xmax = new_target[index][2]
new_ymax = new_target[index][3]
xmin = bndbox.find('xmin')
xmin.text = str(new_xmin)
ymin = bndbox.find('ymin')
ymin.text = str(new_ymin)
xmax = bndbox.find('xmax')
xmax.text = str(new_xmax)
ymax = bndbox.find('ymax')
ymax.text = str(new_ymax)
if xmin.text=="0" and ymin.text=="0" and xmax.text=="0" and ymax.text=="0":
xmlroot.remove(object) # 由于在后续设置中将超出五个裁剪完成图像范围的目标边界框的四个参数全部设为0,此处将参数为0的object删除掉,剩余满足裁剪范围内的目标边界框
index = index + 1
tree.write(os.path.join(saveroot, str(id + '.xml')))
# 遍历指定目录,显示目录下的所有文件名
def CropImage4File(filepath, destpath):
pathDir = os.listdir(filepath) # 列出文件路径中的所有路径或文件
for allDir in pathDir:
child = os.path.join(filepath, allDir)
dest = os.path.join(destpath, allDir)
if os.path.isfile(child):
image = cv2.imread(child)
sp = image.shape # 获取图像形状:返回【行数值,列数值】列表
sz1 = sp[0] # 图像的高度(行 范围)
sz2 = sp[1] # 图像的宽度(列 范围)
# sz3 = sp[2] #像素值由【RGB】三原色组成
# 你想对文件的操作
n, s = os.path.splitext(dest)
cropImg = image
cv2.imwrite(n + s, cropImg) # 写入图像路径
a1 = 0 # x start
b1 = int(sz1 / 2) # x end
c1 = 0 # y start
d1 = int(sz2 / 2) # y end
cropImg1 = image[a1:b1, c1:d1] # 裁剪图像
cv2.imwrite(n+'-1'+s, cropImg1) # 写入图像路径
a2 = 0 # x start
b2 = int(sz1 / 2) # x end
c2 = int(sz2 / 2) # y start
d2 = int(sz2) # y end
cropImg2 = image[a2:b2, c2:d2] # 裁剪图像
cv2.imwrite(n+'-2'+s, cropImg2) # 写入图像路径
a3 = int(sz1 / 2) # x start
b3 = int(sz1) # x end
c3 = 0 # y start
d3 = int(sz2 / 2) # y end
cropImg3 = image[a3:b3, c3:d3] # 裁剪图像
cv2.imwrite(n+'-3'+s, cropImg3) # 写入图像路径
a4 = int(sz1 / 2) # x start
b4 = int(sz1) # x end
c4 = int(sz2 / 2) # y start
d4 = int(sz2) # y end
cropImg4 = image[a4:b4, c4:d4] # 裁剪图像
cv2.imwrite(n+'-4'+s, cropImg4) # 写入图像路径
a5 = int(sz1 / 4) # x start
b5 = int(sz1*3/4) # x end
c5 = int(sz2 / 4) # y start
d5 = int(sz2*3/4) # y end
cropImg5 = image[a5:b5, c5:d5] # 裁剪图像
cv2.imwrite(n + '-5' + s, cropImg5) # 写入图像路径
if __name__ == '__main__':
filepath = '../data/chap/chap2/datasets/03/' # 源图像
XML_DIR = '../data/chap/chap2/datasets/03-1/' # 源xml
destpath = '../data/chap/chap2/datasets/03-2/' # 保存裁剪图片的文件夹地址
AUG_XML_DIR = '../data/chap/chap2/datasets/03-3/' # 保存xml地址
new_bndbox_list = []
for name in tqdm(os.listdir(XML_DIR), desc='Processing'):
bndbox = read_xml_annotation(XML_DIR, name) # 一张图片的边界框信息
og_xml = open(os.path.join(XML_DIR, name))
tree = ET.parse(og_xml)
elem = tree.find('filename')
elem.text = (name[:-4] + '.jpg') # 每一次将xml文件中的filename节点信息进行更改
tree.write(os.path.join(AUG_XML_DIR, name))
img = Image.open(os.path.join(filepath, name[:-4] + '.jpg'))
img = np.asarray(img)
i_h, i_w, _= img.shape
# 分别依据左上右上左下右下和中心位置对满足要求的目标边界框进行分别提取,并保存在各自对应的xml文件中,若目标边界框在裁剪图片边缘上,则舍弃
for i in range(len(bndbox)):
x1=bndbox[i][0]
y1=bndbox[i][1]
x2=bndbox[i][2]
y2=bndbox[i][3]
n_x1 = 0
n_y1 = 0
n_x2 = 0
n_y2 = 0
if x1 < i_w/2 and x2 < i_w/2 and y1 < i_h/2 and y2 < i_h/2:
n_x1 = x1
n_y1 = y1
n_x2 = x2
n_y2 = y2
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
epoch=1
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
# 将new_bndbox_list清零
new_bndbox_list = []
for i in range(len(bndbox)):
x1=bndbox[i][0]
y1=bndbox[i][1]
x2=bndbox[i][2]
y2=bndbox[i][3]
n_x1 = 0
n_y1 = 0
n_x2 = 0
n_y2 = 0
if x1 > i_w/2 and x2 > i_w/2 and y1 < i_h/2 and y2 < i_h/2:
n_x1 = x1
n_y1 = y1
n_x2 = x2
n_y2 = y2
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
epoch=2
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
# 将new_bndbox_list清零
new_bndbox_list = []
for i in range(len(bndbox)):
x1=bndbox[i][0]
y1=bndbox[i][1]
x2=bndbox[i][2]
y2=bndbox[i][3]
n_x1 = 0
n_y1 = 0
n_x2 = 0
n_y2 = 0
if x1 < i_w/2 and x2 < i_w/2 and y1 > i_h/2 and y2 > i_h/2:
n_x1 = x1
n_y1 = y1
n_x2 = x2
n_y2 = y2
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
epoch=3
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
# 将new_bndbox_list清零
new_bndbox_list = []
for i in range(len(bndbox)):
x1=bndbox[i][0]
y1=bndbox[i][1]
x2=bndbox[i][2]
y2=bndbox[i][3]
n_x1 = 0
n_y1 = 0
n_x2 = 0
n_y2 = 0
if x1 > i_w/2 and x2 > i_w/2 and y1 > i_h/2 and y2 > i_h/2:
n_x1 = x1
n_y1 = y1
n_x2 = x2
n_y2 = y2
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
epoch=4
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
# 将new_bndbox_list清零
new_bndbox_list = []
for i in range(len(bndbox)):
x1=bndbox[i][0]
y1=bndbox[i][1]
x2=bndbox[i][2]
y2=bndbox[i][3]
n_x1 = 0
n_y1 = 0
n_x2 = 0
n_y2 = 0
if x1 > i_w/4 and x2 > i_w/4 and y1 > i_h/4 and y2 > i_h/4 and x1 < i_w*3/4 and x2 < i_w*3/4 and y1 < i_h*3/4 and y2 < i_h*3/4:
n_x1 = x1
n_y1 = y1
n_x2 = x2
n_y2 = y2
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
epoch=5
# 存储变化后的XML
change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
CropImage4File(filepath, destpath) # 裁剪图片保存
由于在csdn中并没有找寻到相关裁剪五块图像并使xml文档目标边界框信息跟随裁剪变化的文章,所以研究了一下相关数据扩增的内容,本代码可以完成对目标图片的比例裁剪,以及相应的xml裁剪
效果如下:
通过验证每个部分的边界框数量可以知道本代码是可行的