1 前言
在深度学习的模型训练当中,数据增强可以很好地提高模型地泛化能力和鲁棒性,通过对原有数据进行一定的变换,从而获得更多的数据集,以此来提高模型地表现。数据增强可以通过多种方式实现,比如旋转、缩放、裁剪、变形、镜像、噪声等操作。
由于在进行数据处理时,没有找到合适的 Python 库能够实现将原数据翻转生成新的数据集,并将新数据集与原数据集合并生成新的数据集。在此情形下,本人定义了函数 img_aug() 来实现这一功能。
代码中主要包含以下实现步骤:
- 根据原数据集左右翻转生成新的数据集
- 将新数据集中的图片根据原数据集图片名称统一重命名
- 合并原数据集和新数据集
- 注:进行操作前,需要先创建好相关文件夹 (因为我的目录判空代码写的有点问题……)
2 数据增强
在数据增强这一部分中,代码核心为:imgaug 库,下面是其的简要介绍:
imgaug 是一个用于图像增强的 Python 库,可用于在训练神经网络时生成更多的数据。它包含了多种用于图像增强的方法,包括平移、旋转、缩放、裁剪、翻转、变形等。同时,它还支持随机选择、随机组合这些增强方法,从而生成更多样化的图像数据。
在本文中,使用了 imgaug 库中的 iaa.Sequential 类,用于将多个图像增强操作组合成一个序列,并按顺序依次应用这些操作。下面是iaa.Sequential中一些常用的操作:
-
iaa.Fliplr ( p ):对图像进行水平翻转,p表示翻转的概率。
-
iaa.Flipud( p ):对图像进行垂直翻转,p表示翻转的概率。
-
iaa.Crop( px, keep_size=False):对图像进行随机裁剪,px表示要裁剪的像素数,keep_size表示是否保持裁剪后的大小不变。
-
iaa.Resize({“height”: h, “width”: w}):对图像进行缩放,将其调整为指定的高度和宽度。
-
iaa.Rotate(degrees, fit_output=False):对图像进行旋转,degrees表示旋转的角度,fit_output表示是否调整输出图像的大小以适应旋转后的图像。
-
iaa.Sometimes(p, augmenter):有p的概率应用增强器augmenter,否则不应用任何增强。
-
iaa.GaussianBlur(sigma):对图像进行高斯模糊,sigma表示高斯核的标准差。
-
iaa.Add(value, per_channel=False):对图像进行亮度调整,value表示要添加的值,per_channel表示是否对每个颜色通道应用不同的值。
-
iaa.SaltAndPepper§:对图像添加椒盐噪声,p表示添加噪声的概率。
2.1 生成目录列表
由于数据集文件是一个多级目录,所以在下面的代码中,需要先获取照片所对应的上一级目录列表,并将其存放在列表当中,实现代码如下:
import os
input_path = r"G:\Desktop\NJUPT\GitHub\365_pytorch_learn\demo\data\ori_data" # 原数据集路径
output_path = r"G:\Desktop\NJUPT\GitHub\365_pytorch_learn\demo\data\res_data" # 反转后存放数据集的路径
# 路径地址最好使用绝对路径
input_dirs = [d for d in os.listdir(input_path) if os.path.isdir(os.path.join(input_path, d))]
output_dirs = [d for d in os.listdir(output_path) if os.path.isdir(os.path.join(output_path, d))]
for i in range(len(input_dirs)):
input_dirs[i] = input_path + "\\"+ input_dirs[i]
for i in range(len(output_dirs)):
output_dirs[i] = output_path + "\\"+ output_dirs[i]
print(input_dirs)
print(output_dirs)
数据结果如下:
2.2 数据增强
接下来使用 imgaug 库对 input_dirs 目录中的一组图像进行数据增强,并将增强后的图像保存在 output_dirs 目录中。
该代码使用垂直翻转的增强方式,概率为50%。代码遍历 input_dirs 目录中的每个子目录中的每个图像文件,检查文件是否为 JPEG 或 PNG 图像,打开图像文件并将其转换为 NumPy 数组。
然后,将增强序列应用于图像数组,将增强后的数组转换回 PIL Image 对象,并将增强后的图像保存到 output_dirs 目录中。
此外,在实际操作中,发现数据集中存在一些图片为 RGBA 格式图片,所以代码代码还检查原始图像是否为RGBA格式,并在应用增强之前将其转换为RGB格式。
实现代码如下:
import imgaug.augmenters as iaa
import os
from PIL import Image
import numpy as np
seq = iaa.Sequential([
iaa.Flipud(1.0), # 按垂直轴随机翻转50%的图片
])
len_class = len(input_dirs) # 获取文件夹中子目录长度
for i in range(len_class):
for file_name in os.listdir(input_dirs[i]):
if file_name.endswith('.jpg') or file_name.endswith('.jpeg') or file_name.endswith('.png'):
file_path = os.path.join(input_dirs[i], file_name)
img = Image.open(file_path)
if img.mode == 'RGBA':
img = img.convert('RGB')
img_arr = np.array(img)
img_aug = seq(images=img_arr)
img_aug = Image.fromarray(img_aug)
if img_aug.mode == 'RGBA':
img_aug = img_aug.convert('RGB')
output_file_path = os.path.join(output_dirs[i], file_name)
img_aug.save(output_file_path)
3 图片重命名
在上面代码所生成的新数据集中的图片名称并没有改变,因此需要将新数据集中的图片进行重命名,以便后续合并数据集以及模型训练。
下面的代码作用为,统计原数据及中的图片个数,以便获取新数据中起始图片的序列号。
import os
import glob
folder_path = r'G:\Desktop\NJUPT\GitHub\365_pytorch_learn\demo\data\res_data' # 修改为你的文件夹路径
counters = {} # 定义一个空字典
for folder_name in os.listdir(folder_path):
folder = os.path.join(folder_path, folder_name)
if os.path.isdir(folder):
file_types = ['*.jpg', '*.jpeg', '*.png'] # 文件类型可以根据实际情况修改
total_files = 0
for file_type in file_types:
total_files += len(glob.glob(os.path.join(folder, file_type)))
counters[folder_name] = total_files # 将文件夹名和图片数量添加到字典中
# 使用字典推导式将字典中的每个键值加1
counters = {key: value + 1 for key, value in counters.items()}
print(counters)
输出结果为:{‘cloudy’: 301, ‘rain’: 431}
接下来遍历每个子文件夹中的图片,并对其重命名操作,实现代码如下:
# 将一个文件夹下的所有图片文件按照特定的命名规则重命名
import os
import re
# 遍历每个子文件夹
for subfolder in os.listdir(folder_path):
subfolder_path = os.path.join(folder_path, subfolder)
# 如果是文件夹,则遍历文件夹中所有文件
if os.path.isdir(subfolder_path):
# 获取当前子文件夹的起始计数器
counter = counters[subfolder]
for name in os.listdir(subfolder_path):
# 如果是图片文件,则解析文件名并重命名文件
if name.endswith('.jpg') or name.endswith('.jpeg') or name.endswith('.png'):
# 解析文件名中的关键信息
pattern = r'(\w+)_(\d+)\.(jpg|jpeg|png)'
match = re.match(pattern, name)
if match:
keyword, number, extension = match.groups()
else:
keyword = subfolder
number = counter
extension = os.path.splitext(name)[1][1:]
# 构造新的文件名
new_name = keyword + str(number) + '.' + extension
# 使用os.rename函数重命名文件
os.rename(os.path.join(subfolder_path, name), os.path.join(subfolder_path, new_name))
# 更新计数器变量
counter += 1
4 合并两个文件中的图片
此部分没有废话,直接上代码:
# 将生成的新照片与原照片合并
import os
import shutil
for i in range(len_class):
# 如果合并文件夹不存在,创建一个新的文件夹
if not os.path.exists(new_data_dirs[i]):
os.mkdir(new_data_dirs[i])
# 遍历文件夹1中的所有文件,并复制到合并文件夹中
for filename in os.listdir(input_dirs[i]):
if filename.endswith(".jpg"): # 只复制图片文件
src_path = os.path.join(input_dirs[i], filename)
dst_path = os.path.join(new_data_dirs[i], filename)
shutil.copyfile(src_path, dst_path)
# 遍历文件夹2中的所有文件,并复制到合并文件夹中
for filename in os.listdir(output_dirs[i]):
if filename.endswith(".jpg"): # 只复制图片文件
src_path = os.path.join(output_dirs[i], filename)
dst_path = os.path.join(new_data_dirs[i], filename)
shutil.copyfile(src_path, dst_path)
print("图片已成功合并到新文件夹!")
5 完整代码
此代码实现了将数据集左右翻转并合并生成新的数据集
待改进:
-
文件夹的判空操作,本代码在处理数据之前,应该先将文件夹创建好,架构如下(见代码注释部分)
-
此代码只能实现图片左右翻转,如需实现上下反转以及其他的数据增强操作,需重写 imgaug() 函数下的 seq 部分(建议查看函数源码) (line50)
"""
此代码实现了将数据集左右翻转并合并生成新的数据集
待改进:
* 文件夹的判空操作,本代码在处理数据之前,应该先将文件夹创建好,架构如下
|-data
|-ori_data
|-res_data
|-new_data
* 此代码只能实现图片左右翻转,如需实现上下反转以及其他的数据增强操作,需重写 imgaug() 函数下的 seq 部分(建议查看函数源码)(line48)
"""
import imgaug.augmenters as iaa
import os
from PIL import Image
import numpy as np
import re
import shutil
import glob
def data_aug(input_path, output_path, new_data_path):
"""主函数,用于调用其他函数"""
if not os.path.exists(input_path):
os.mkdir(input_path)
if not os.path.exists(output_path):
os.mkdir(output_path)
if not os.path.exists(new_data_path):
os.mkdir(new_data_path)
imgaug(input_path, output_path)
rename_img(output_path)
merge_files(new_data_path, input_path, output_path)
def imgaug(input_path, output_path):
"""获取原图片文件路径和翻转后生成图片路径"""
input_dirs = [d for d in os.listdir(input_path) if os.path.isdir(os.path.join(input_path, d))]
output_dirs = [d for d in os.listdir(output_path) if os.path.isdir(os.path.join(output_path, d))]
for i in range(len(input_dirs)):
input_dirs[i] = input_path + "\\" + input_dirs[i]
for i in range(len(output_dirs)):
output_dirs[i] = output_path + "\\" + output_dirs[i]
seq = iaa.Sequential([
iaa.Flipud(1.0), # 按垂直轴随机翻转50%的图片
])
len_class = len(input_dirs) # 获取文件夹中子目录长度
for i in range(len_class):
for file_name in os.listdir(input_dirs[i]):
if file_name.endswith('.jpg') or file_name.endswith('.jpeg') or file_name.endswith('.png'):
file_path = os.path.join(input_dirs[i], file_name)
img = Image.open(file_path)
if img.mode == 'RGBA':
img = img.convert('RGB')
img_arr = np.array(img)
img_aug = seq(images=img_arr)
img_aug = Image.fromarray(img_aug)
if img_aug.mode == 'RGBA':
img_aug = img_aug.convert('RGB')
output_file_path = os.path.join(output_dirs[i], file_name)
img_aug.save(output_file_path)
def rename_img(folder_path):
"""将生成的新文件夹中图片按照顺序重命名"""
counters = {} # 定义一个空字典
for folder_name in os.listdir(folder_path):
folder = os.path.join(folder_path, folder_name)
if os.path.isdir(folder):
file_types = ['*.jpg', '*.jpeg', '*.png'] # 文件类型可以根据实际情况修改
total_files = 0
for file_type in file_types:
total_files += len(glob.glob(os.path.join(folder, file_type)))
counters[folder_name] = total_files # 将文件夹名和图片数量添加到字典中
# 使用字典推导式将字典中的每个键值加1
counters = {key: value + 1 for key, value in counters.items()}
# 遍历每个子文件夹
for subfolder in os.listdir(folder_path):
subfolder_path = os.path.join(folder_path, subfolder)
# 如果是文件夹,则遍历文件夹中所有文件
if os.path.isdir(subfolder_path):
# 获取当前子文件夹的起始计数器
counter = counters[subfolder]
for name in os.listdir(subfolder_path):
# 如果是图片文件,则解析文件名并重命名文件
if name.endswith('.jpg') or name.endswith('.jpeg') or name.endswith('.png'):
# 解析文件名中的关键信息
pattern = r'(\w+)_(\d+)\.(jpg|jpeg|png)'
match = re.match(pattern, name)
if match:
keyword, number, extension = match.groups()
else:
keyword = subfolder
number = counter
extension = os.path.splitext(name)[1][1:]
# 构造新的文件名
new_name = keyword + str(number) + '.' + extension
# 使用os.rename函数重命名文件
os.rename(os.path.join(subfolder_path, name), os.path.join(subfolder_path, new_name))
# 更新计数器变量
counter += 1
def merge_files(new_data_path, input_path, output_path):
"""合并原数据集和新数据集"""
new_data_dirs = [d for d in os.listdir(new_data_path) if os.path.isdir(os.path.join(new_data_path, d))]
for i in range(len(new_data_dirs)):
new_data_dirs[i] = new_data_path + "\\" + new_data_dirs[i]
len_class = len(new_data_dirs)
input_dirs = [d for d in os.listdir(input_path) if os.path.isdir(os.path.join(input_path, d))]
output_dirs = [d for d in os.listdir(output_path) if os.path.isdir(os.path.join(output_path, d))]
for i in range(len(input_dirs)):
input_dirs[i] = input_path + "\\" + input_dirs[i]
for i in range(len(output_dirs)):
output_dirs[i] = output_path + "\\" + output_dirs[i]
for i in range(len_class):
# 如果合并文件夹不存在,创建一个新的文件夹
if not os.path.exists(new_data_dirs[i]):
os.mkdir(new_data_dirs[i])
# 遍历文件夹1中的所有文件,并复制到合并文件夹中
for filename in os.listdir(input_dirs[i]):
if filename.endswith(".jpg"): # 只复制图片文件
src_path = os.path.join(input_dirs[i], filename)
dst_path = os.path.join(new_data_dirs[i], filename)
shutil.copyfile(src_path, dst_path)
# 遍历文件夹2中的所有文件,并复制到合并文件夹中
for filename in os.listdir(output_dirs[i]):
if filename.endswith(".jpg"): # 只复制图片文件
src_path = os.path.join(output_dirs[i], filename)
dst_path = os.path.join(new_data_dirs[i], filename)
shutil.copyfile(src_path, dst_path)
print("图片已成功合并到新文件夹!")
if __name__ == "__main__":
input_path = r"G:\Desktop\NJUPT\GitHub\365_pytorch_learn\demo\data\ori_data" # 原数据集路径
output_path = r"G:\Desktop\NJUPT\GitHub\365_pytorch_learn\demo\data\res_data" # 反转后存放数据集的路径
new_data_path = r"G:\Desktop\NJUPT\GitHub\365_pytorch_learn\demo\new_data" # 最终两个数据集合并后的路径
data_aug(input_path, output_path, new_data_path)