同时翻转图片和xml扩增数据集

1、水平翻转:

# -*- coding: utf-8 -*-
"""
Created on Mon Jul 26 11:31:01 2021

@author: admin
"""
import glob
from PIL import Image
import xml.etree.ElementTree as ET
import xml.dom.minidom
from tqdm import tqdm
import os

def flip_xml(w, cs, newname, xml_file, new_xml):
    tree = ET.parse(xml_file)
    objs = tree.findall('object')
    for ix, obj in enumerate(objs):
        if obj.find('name').text in cs:
            print(xml_file)
            obj_new = obj.find('bndbox')
            xmin = str( w - int(obj_new.find('xmin').text))
            xmax = str( w - int(obj_new.find('xmax').text))
            obj_new.find('xmin').text = xmax
            obj_new.find('xmax').text = xmin
            tree.write(new_xml)
    
    dom = xml.dom.minidom.parse(new_xml)
    root = dom.documentElement
 
    # obtain the filename label pair and give it a new value
    root.getElementsByTagName('filename')[0].firstChild.data = newname + '.jpg'
    root.getElementsByTagName('path')[0].firstChild.data = 'E:\data\\' + newname + '.jpg'
    with open(new_xml,'w') as fh:
        dom.writexml(fh)


if __name__ == '__main__':
    src=r'data'
    dst = 'data_flip'
    cs = ['flower', 'person']
    xml_list= glob.glob(src + '/*.xml')
    for file_name in tqdm(xml_list):
        xml_name = file_name.split('\\')[-1]#linux
        name = xml_name.split('.')[0]
        newname = name + '_fh'
        
        jpg_file = os.path.join(src, name + '.jpg')
        new_jpg = os.path.join(dst, newname + '.jpg')
        im = Image.open(jpg_file)
        (w,h) = im.size
        im_h = im.transpose(Image.FLIP_LEFT_RIGHT)
        im_h.save(new_jpg)
    
        new_xml = os.path.join(dst, newname + '.xml') 
        flip_xml(w, cs, newname, file_name, new_xml)
        

2、任意角度旋转:

# 读取图片,读取xml文件
# 解析xml文件获取图片的大小,xml中的多组坐标
# 翻转图片,根据图片翻转的角度,获取新的图片大小和坐标点
# 将坐标信息和图片大小在原有的xml中修改参数保存到新的xml文件中
# 创建新的文件夹,将图片和xml文件放在同一文件夹中(一组图片和xml文件命名相同)
 
import os
import numpy as np
import cv2
import math
import xml.etree.ElementTree as ET
 
#按角度翻转图片
def rotate_img(src, angle, scale = 1):
	width = src.shape[1]
	height = src.shape[0]
	# 角度变弧度
	re_angle = np.deg2rad(angle)
	# 计算新图片的高度和宽度
	new_width = (abs(np.sin(re_angle) * height) + abs(np.cos(re_angle) * width)) * scale
	new_height = (abs(np.cos(re_angle) * height) + abs(np.sin(re_angle) * width)) * scale
 
	rotate_matrix = cv2.getRotationMatrix2D((new_width * 0.5, new_height * 0.5), angle, scale)
	rotate_move = np.dot(rotate_matrix, np.array([(new_width - width) * 0.5, (new_height - height) * 0.5, 0]))
 
	# update translation
	rotate_matrix[0, 2] += rotate_move[0]
	rotate_matrix[1, 2] += rotate_move[1]
 
	dst = cv2.warpAffine(img, rotate_matrix, (int(math.ceil(new_width)), int(math.ceil(new_height))), flags=cv2.INTER_LANCZOS4)
	return dst
 
# 翻转后的xml文件信息
def rotate_xml(src, xmin, ymin, xmax, ymax, angle, scale = 1):
	width = src.shape[1]
	height = src.shape[0]
	re_angle = np.deg2rad(angle)
	new_width = (abs(np.sin(re_angle) * height) + abs(np.cos(re_angle) * width)) * scale
	new_height = (abs(np.cos(re_angle) * height) + abs(np.sin(re_angle) * width)) * scale
	rotate_matrix = cv2.getRotationMatrix2D((new_width * 0.5, new_height * 0.5), angle, scale)
	rotate_move = np.dot(rotate_matrix, np.array([(new_width - width) * 0.5, (new_height - height) * 0.5, 0]))
	rotate_matrix[0, 2] += rotate_move[0]
	rotate_matrix[1, 2] += rotate_move[1]
	# 获取原始矩形的四个中点,然后将这四个点转换到旋转后的坐标系下
	point1 = np.dot(rotate_matrix, np.array([(xmin + xmax) / 2, ymin, 1]))
	point2 = np.dot(rotate_matrix, np.array([xmax, (ymin + ymax) / 2, 1]))
	point3 = np.dot(rotate_matrix, np.array([(xmin + xmax) / 2, ymax, 1]))
	point4 = np.dot(rotate_matrix, np.array([xmin, (ymin + ymax) / 2, 1]))
	concat = np.vstack((point1, point2, point3, point4))  # 合并np.array
	# 改变array类型
	concat = concat.astype(np.int32)
	rx, ry, rw, rh = cv2.boundingRect(concat)		#rx,ry,为新的外接框左上角坐标,rw为框宽度,rh为高度
	new_xmin = rx
	new_ymin = ry
	new_xmax = rx + rw
	new_ymax = ry + rh
 
	return new_xmin, new_ymin, new_xmax, new_ymax
 
if __name__ == '__main__':
 
	file_path = './data/' #输入路径
	rotated_path = './data_flip/'	 #翻转后的路径
 
	# 自定义翻转角度
	for angle in (10,20,30,60,90,180,270,300,330,340,350):
		for file in os.listdir(file_path):
			if file.endswith('.jpg'):
				a,b = os.path.splitext(file)
				img = cv2.imread(file_path + a + '.jpg')
				rotated_img = rotate_img(img, angle)
				width_d = rotated_img.shape[1]
				height_d = rotated_img.shape[0]
				cv2.imwrite(rotated_path + a + '_' + str(angle) + 'd.jpg', rotated_img)
				print(str(file) + ' ' + 'has been rotated for' + str(angle) + '°')
			if file.endswith('.xml'):
				src = cv2.imread(file_path + a + '.jpg')
				tree = ET.parse(file_path + a + '.xml')
				root = tree.getroot()
				root.find('folder').text = 'flower'
				root.find('filename').text = a + '_' + str(angle) + 'd.jpg'
				root.find('path').text = 'E:\data\\flower\\' + a + '_' + str(angle) + 'd.jpg'
				root.find("size").find('width').text = str(width_d)
				root.find("size").find('height').text = str(height_d)
				# 修改xml中的标签坐标信息
				for box in root.iter('bndbox'):
					xmin = float(box.find('xmin').text)
					ymin = float(box.find('ymin').text)
					xmax = float(box.find('xmax').text)
					ymax = float(box.find('ymax').text)
					new_xmin, new_ymin, new_xmax, new_ymax = rotate_xml(src, xmin, ymin, xmax, ymax, angle)
					box.find('xmin').text = str(max(new_xmin, 0))
					box.find('ymin').text = str(max(new_ymin, 0))
					box.find('xmax').text = str(min(new_xmax, width_d))
					box.find('ymax').text = str(min(new_ymax, height_d))
				tree.write(rotated_path + a + '_' + str(angle) + 'd.xml')
				print(str(file) + ' ' + 'has been rotated for ' + str(angle) + '°')
 
	print("-----------------------------------")
	print("Successful!")

参考文章:https://blog.csdn.net/qq_41192383/article/details/95140597

                 https://blog.csdn.net/weixin_42273095/article/details/108779096                                       

  • 0
    点赞
  • 16
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值