数据增强只针对一张图片的话效率太低了,因此以下代码都是可以对文件夹中的图片集体进行数据增强。
1 平移
import cv2 as cv
import os
import numpy as np
original_path = '../images' # 原图文件夹路径
output_path = '../output' # 输出文件夹名称
distance = 50 # 单位是像素
def image_translate(mode,distance,original_path=None,output_path=None):
# 遍历文件夹中的所有图片
for name in os.listdir(original_path):
if name.endswith('.jpg') or name.endswith('.png') or name.endswith('.jpeg'):
original_image = cv.imread(original_path + '/' + name)
savepath = (output_path + '/' + name[:name.index('.')] + '_' + str(mode) + '_' + str(distance) + name[name.index('.'):])
# print('savepath:', savepath)
if mode == 'Up':
M = np.float32([[1, 0, 0], [0, 1, -distance]])
shifted = cv.warpAffine(original_image, M, (original_image.shape[1], original_image.shape[0]))
cv.imwrite(savepath, shifted)
if mode == 'Down':
M = np.float32([[1, 0, 0], [0, 1, distance]])
shifted = cv.warpAffine(original_image, M, (original_image.shape[1], original_image.shape[0]))
cv.imwrite(savepath, shifted)
if mode == 'Left':
M = np.float32([[1, 0, -distance], [0, 1, 0]])
shifted = cv.warpAffine(original_image, M, (original_image.shape[1], original_image.shape[0]))
cv.imwrite(savepath, shifted)
if mode == 'Right':
M = np.float32([[1, 0, distance], [0, 1, 0]])
shifted = cv.warpAffine(original_image, M, (original_image.shape[1], original_image.shape[0]))
cv.imwrite(savepath, shifted)
image_translate('Up', distance, original_path=original_path, output_path=output_path)
image_translate('Down', distance, original_path=original_path, output_path=output_path)
image_translate('Left', distance, original_path=original_path, output_path=output_path)
image_translate('Right', distance, original_path=original_path, output_path=output_path)
print('finished')
2 翻转
import cv2 as cv
import os
"""
<mode1>:
'single':对单个图像进行操作。'folder':对文件夹内所有图像批量操作。
<mode2>:
'UD':图像上下翻转。'LR':图像左右翻转。'T':图像转置。
------------------------------------------------------------------
<singlepath>:单个图像的路径。
<original_path>:文件夹批量操作,文件夹的路径。
<new_path>:文件夹批量操作,新生成文件存储文件夹的路径。
"""
original_path = '../images' # 原图文件夹路径
output_path = '../output' # 输出文件夹路径
def image_flip(mode1, mode2, singlepath=None, original_path=None, new_path=None):
if new_path is not None and os.path.exists(new_path) is False:
os.makedirs(new_path)
if mode1 == 'folder':
for name in os.listdir(original_path):
original_image = cv.imread(original_path + '/' + name)
savepath = (new_path + '/' + name[:name.index('.')] + '_' + str(mode2) + name[name.index('.'):])
# print('savepath:', savepath)
if mode2 == 'UD':
new_image = cv.flip(original_image, 0) # original_image[::-1] 上下翻转
cv.imwrite(savepath, new_image)
elif mode2 == 'LR':
new_image = cv.flip(original_image, 1) # original_image[:, ::-1] 左右翻转
cv.imwrite(savepath, new_image)
elif mode2 == 'T':
new_image = cv.flip(original_image, -1) # 水平垂直同时翻转
cv.imwrite(savepath, new_image)
else:
print("mode2可选参数为:\n'UD':上下翻转\n'LR':左右翻转\n'T' :转置\n请输入正确参数!!!")
break
elif mode1 == 'single':
original_image = cv.imread(singlepath)
savepath = singlepath[:singlepath.index('.')] + str(mode2) + singlepath[singlepath.index('.'):]
if mode2 == 'UD':
new_image = cv.flip(original_image, 0) # original_image[::-1] 上下翻转
cv.imwrite(savepath, new_image)
if mode2 == 'LR':
new_image = cv.flip(original_image, 1) # original_image[:, ::-1] 左右翻转
cv.imwrite(savepath, new_image)
if mode2 == 'T':
new_image = cv.flip(original_image, -1) # 水平垂直同时翻转
cv.imwrite(savepath, new_image)
else:
print("mode1可选参数为:\n'folder':批量处理\n'single':处理单张\n请输入正确参数!!!")
os.rmdir(new_path)
# 上下翻转
image_flip('folder', 'UD', original_path=original_path, new_path=output_path)
# 左右翻转
image_flip('folder', 'LR', original_path=original_path, new_path=output_path)
# 转置(先上下后左右翻转)
image_flip('folder', 'T', original_path=original_path, new_path=output_path)
print('finished')
3 旋转
import os
from PIL import Image
original_path = '../images' # 原图文件夹路径
output_path = '../output' # 输出文件夹名称
angle = 90 # 旋转角度
def image_rotate(angle, original_path=None, output_path=None):
for name in os.listdir(original_path):
savepath = (output_path + '/' + name[:name.index('.')] + '_rotate_' + str(angle) + name[name.index('.'):])
im = Image.open(original_path + '/' + name)
im = im.convert("RGB") # 把PNG格式转换成的四通道转成RGB的三通道
im_rotate = im.rotate(angle, expand=1) # 逆时针旋转90度,expand=1表示原图直接旋转
# 判断输出文件夹是否已存在,不存在则创建。
folder = os.path.exists(output_path)
if not folder:
os.makedirs(output_path)
# 把旋转后的图片存入输出文件夹
im_rotate.save(savepath)
image_rotate(angle, original_path=original_path, output_path=output_path)
image_rotate(angle + 90, original_path=original_path, output_path=output_path)
image_rotate(angle + 180, original_path=original_path, output_path=output_path)
print('finished')
参考博文:
https://blog.csdn.net/weixin_47370211/article/details/123756361
https://blog.csdn.net/github_36350097/article/details/117885214
批量重命名
import os
image_path = '../images' # 表示需要命名处理的文件夹
def rename(image_path):
filelist = os.listdir(image_path) # 获取文件路径
total_num = len(filelist) # 获取文件长度(个数)
i = 0 # 表示文件的命名是从0开始的
for item in filelist:
if item.endswith('.jpg') or item.endswith('.png') or item.endswith('.jpeg'):
src = os.path.join(os.path.abspath(image_path), item)
dst = os.path.join(os.path.abspath(image_path),
'0' + format(str(i), '0>3s') + '.jpg') # 这种情况下的命名格式为0000.jpg形式,可以自主定义想要的格式
try:
os.rename(src, dst)
print('converting %s to %s ...' % (src, dst))
i = i + 1
except:
continue
print('total %d to rename & converted %d jpgs' % (total_num, i))
rename(image_path=image_path)
划分训练集和数据集
import os
import random
import shutil
# 示例数据路径
data_dir = '../custom_dataset/images' # 替换为你的数据路径
# 创建训练集和验证集文件夹
train_dir = '../custom_dataset/train'
val_dir = '../custom_dataset/val'
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
# 设置随机数种子,以确保每次运行代码时划分结果一致
random.seed(42)
# 定义划分比例,这里是80%的训练集和20%的验证集
split_ratio = 0.8
# 获取数据集文件列表
file_list = os.listdir(data_dir)
# 随机打乱文件列表
random.shuffle(file_list)
# 计算划分的索引
split_index = int(len(file_list) * split_ratio)
# 划分文件列表
train_files = file_list[:split_index]
val_files = file_list[split_index:]
# 将训练集文件复制到训练集文件夹
for file in train_files:
src_path = os.path.join(data_dir, file)
dest_path = os.path.join(train_dir, file)
shutil.copy(src_path, dest_path)
# 将验证集文件复制到验证集文件夹
for file in val_files:
src_path = os.path.join(data_dir, file)
dest_path = os.path.join(val_dir, file)
shutil.copy(src_path, dest_path)
print("训练集文件夹:", train_dir)
print("验证集文件夹:", val_dir)