自制数据集之数据增强（翻转、平移、旋转）

数据增强只针对一张图片的话效率太低了，因此以下代码都是可以对文件夹中的图片集体进行数据增强。

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)