[数据集处理]数据增强1

工欲善其事，必先利其器

[数据集处理]数据增强2 参考了 yolo v4 马赛克数据增强方法，进行img\xml的四合一

记录一个标注过的图片水平翻转，对应的xml标注信息也同步的小脚本

import cv2
import os
import sys
import re
import xml.etree.ElementTree as ET
from PIL import Image


imgreadpath = './img/'
imgwritepath  = './f_img/'

xmlreadpath = './xml/'
xmlwritepath = './f_xml/'


def flitimg(imgname):
	image = cv2.imread(imgname)
	name = imgname.split('/')[-1].split('.')[-2]
	image_f = cv2.flip(image, 1)  #1：水平翻转
	cv2.imwrite(imgwritepath + 'f' + name + '.jpg',image_f)

def flitxml(xmlname):

	bwidth = ''
	bheight = ''
	bdepth = ''

	text=open(xmlname).read()
	f_test = open(xmlwritepath + 'f' + xmlname.split("/")[-1].split('.')[-2] + '.xml', 'w')

	text=re.sub(u"[\x00-\x08\x0b-\x0c\x0e-\x1f]+",u"",text)
	root=ET.fromstring(text)

	if root.findtext('folder') == ' ':
		print("no folder filename:",filename)
		print('A')
	else:
		folder = root.findtext('folder')
	
	if root.findtext('filename') == ' ':
		print("no filename filename:",filename)
		print('B')
	else:
		filename = root.findtext('filename')

	# if root.findall('size') == []:
	# 	print("no size filename:",filename)
	# 	print('C')
	# else:
	# 	size = root.findall('size')
	# 	print(size)
	# 	size = size[0]
	# 	bwidth = size.findtext('width')
	# 	bheight = size.findtext('height')
	# 	bdepth = size.findtext('depth')
	img = Image.open(imgwritepath + 'f' + xmlname.split("/")[-1].split('.')[-2] + '.jpg')
	bwidth, bheight = img.size
	bdepth  =3

	f_test.write('<annotation>\n')
	f_test.write('	<folder>' + folder + '</folder>\n')
	f_test.write('	<filename>' + filename + '</filename>\n')
	f_test.write('	<size>\n')
	f_test.write('		<width>' + str(bwidth) + '</width>\n')
	f_test.write('		<height>' + str(bheight) + '</height>\n')
	f_test.write('		<depth>' + str(bdepth) + '</depth>\n')
	f_test.write('	</size>\n')

	if root.findall('object') == []:
		print("no object filename:",filename)
	else:
		for object in root.findall('object'):
			label = object.findtext('name')
			x1 = object.findtext('bndbox/xmin')
			y1 = object.findtext('bndbox/ymin')
			x2 = object.findtext('bndbox/xmax')
			y2 = object.findtext('bndbox/ymax')

			f_test.write('	<object>\n')
			f_test.write('		<name>' + label + '</name>\n')
			f_test.write('		<bndbox>\n')
			f_test.write('			<xmin>' + str(int(bwidth) - int(x1) - (int(x2) - int(x1))) + '</xmin>\n')
			f_test.write('			<ymin>' + y1 + '</ymin>\n')
			f_test.write('			<xmax>' + str(int(bwidth) - int(x2) + (int(x2) - int(x1))) + '</xmax>\n')
			f_test.write('			<ymax>' + y2 + '</ymax>\n')
			f_test.write('		</bndbox>\n')
			f_test.write('	</object>\n')

	f_test.write('</annotation>\n')
	f_test.close()





if __name__ == '__main__':

	imgnames = os.listdir(imgreadpath)
	for imgname in imgnames:
	    imgname = imgreadpath + imgname
		flitimg(imgname)

	xmlnames = os.listdir(xmlreadpath)
	for xmlname in xmlnames:
		xmlname = xmlreadpath + xmlname
		flitxml(xmlname)

代码粘上CSDN后缩进似乎有点问题。。。

原创小脚本，希望能对你有所帮助！

(•̀ᴗ•́)و ̑̑点赞