数据增强python代码

一 批量处理数据集中的图片

查看数据集中每个文件夹中的图片个数，并裁剪 ，灰度化图片

from skimage import data_dir, io, transform, color
from skimage.transform import rotate
import numpy as np
import skimage.io as io
# 读取文件夹
data_dir='F:/project/fzq/data/train/baishao'
str=data_dir + '/*.jpg'
# 显示文件夹中有多少张图 # 混合图片str='d:/pic/*.jpg:d:/pic/*.png'
coll = io.ImageCollection(str)
print(len(coll))

# 显示第i张图片
# io.imshow(coll[10])

#---------------------转化为灰度图片-----------------------
def convert_gray(f):
    rgb = io.imread(f)  # 依次读取rgb图片
    cropped = rgb[200:(rgb.shape[0] - 200), 200:(rgb.shape[1] - 200)]
    # gray = color.rgb2gray(rgb)  # 将rgb图片转换成灰度图
    # image_rotated = rotate(gray, angle=45, resize=True)
    #dst = transform.resize(gray)  #, (256, 256))  # 将灰度图片大小转换为256*256
    return cropped


coll = io.ImageCollection(str, load_func=convert_gray)
for i in range(len(coll)):
    io.imsave('F:/project/fzq/data/train/baishao/' + 'caijian.' +np.str(i) + '.jpg', coll[i])  # 循环保存图片

二 对于某一种图片进行数据处理

擦除+缩小+裁剪+旋转+翻转+颜色（灰度化，饱和度，对比度）+噪声

from PIL import Image
import numpy as np
import torch
from torchvision.transforms import RandomErasing

img = Image.open('img.JPG')
print(type(img))  # 图像的格式
print(np.array(img).shape) # 图像的（长，宽，通道）

# img.show()  # 原图

# ------------- 擦除----------------

img_array = np.array(img).transpose(2, 0, 1)  # transpose (H, W, C) -> (C, H, W)
img_tensor = torch.from_numpy(img_array)

# random_erasing = RandomErasing(
#     p=1.0,  # 概率值，执行该操作的概率，默认为 0.5
#     scale=(0.02, 0.33),  # 按均匀分布概率抽样，遮挡区域的面积 = image * scale
#     ratio=(0.3, 3.3),  # 遮挡区域的宽高比，按均匀分布概率抽样
#     value='123',  # 遮挡区域的像素值，(R, G, B) or (Gray)；传入字符串表示用随机彩色像素填充遮挡区域
#     inplace=False
# )
# 注意，随机遮挡是对 (c, h, w) 形状的 tensor 进行操作，一般在 ToTensor 之后进行
# erased_img_tensor = random_erasing(img_tensor)
# erased_img_array = erased_img_tensor.numpy().transpose(1, 2, 0)  # (C, H, W) -> (H, W, C) 重新变为长宽通道
# erased_img_pil = Image.fromarray(erased_img_array)
# erased_img_pil.show() # 擦除图

#----------------变形-----------------
from torchvision.transforms import Resize

resize = Resize(
    size=(224, 224),  # (height, width)
    interpolation=2  # 插值方法，一般保持默认就好
)
# resize(img).show()

# ----------------裁剪--------------------
from torchvision.transforms import CenterCrop,RandomCrop
from PIL import Image

center_crop = CenterCrop(size=(224, 224))
# center_crop(img).show()

random_crop = RandomCrop(
    size=(224, 224),  # 裁剪后图片的尺寸
    padding=(75, 75),  # (左右填充多少，上下填充多少)
    pad_if_needed=False,  # 当 size 大于原始图片的尺寸时，必须将该参数设置为 True，否则会报错
    fill=(255, 0, 0),  # 同 Pad 对象的参数
    padding_mode='constant'  # 同 Pad 对象的参数
)
# random_crop(img).show()

#--------------------旋转-----------------------
from torchvision.transforms import RandomRotation

random_rotation = RandomRotation(
    degrees=(-45, 45),  # 传入为整数 d 时，旋转角度在 (-d, d)；传入整数元组 (d1, d2)时，角度在 (d1, d2)
    resample=False,  # 旋转之后会要重采样，一般设置为默认就好
    expand=True,  # 传入 True 时，会 padding 图片，保证旋转后四个角的信息不丢失
    center=(0, 0)  # 以什么点为轴旋转，默认为中心，可传入 tuple 更改为任意坐标
)
# 注意，若对不止一张 image 进行处理，在传入 expand=True 时，每张图片的 size 应该要相同，否则报错
# expand=True 是针对以中心为轴旋转所需要 padding 的量的，因此当center=(0, 0)，信息依旧会丢失

# random_rotation(img).show()
#---------------------翻转----------------------------
from torchvision.transforms import RandomHorizontalFlip
image = Image.open('img.JPG')

random_horizontal_flip = RandomHorizontalFlip(p=1.0)  # p 为翻转的概率，默认为 0.5
# random_horizontal_flip(image).show()  # 水平翻转

from torchvision.transforms import RandomVerticalFlip

random_vertical_flip = RandomVerticalFlip(p=1.0) # p 为翻转的概率，默认为 0.5
# random_vertical_flip(image).show() # 上下翻转

#-----------------------颜色----------------------------------
from PIL import Image
from torchvision.transforms import Grayscale,ColorJitter

grayscale = Grayscale(
    num_output_channels=3  # num_output_channels should be either 1 or 3
)
# grayscale(image).show() 灰度化

color_jitter = ColorJitter(
    brightness=(1.5, 1.5),  # 亮度，当传入 a 时，从 [max(0, 1-a), 1+a] 中按均匀分布抽样
                            # 当传入 (a, b) 时，从 [a, b] 中按均匀分布抽样
    contrast=0,  # 对比度，传入格式同 brightness
    saturation=0,  # 饱和度，传入格式同 brightness
    hue=0  # 色相，当传入 a 时，从 [-a, a] 中按均匀分布抽样，注意，0 <= a <= 0.5
           # 当传入 (a, b) 时，从 [a, b] 中按均匀分布抽样，注意，-0.5  <= a <= b <= 0.5
)
# color_jitter(image).show() # 对比度

color_jitter = ColorJitter(
    brightness=0,
    contrast=0,
    saturation=(1.5, 1.5),
    hue=0
)
# color_jitter(image).show() # 饱和度

# ----------------噪声------------------
import random
from PIL import Image
import skimage.io
import skimage.util
import numpy as np


class RandomNoise:
    """按概率为图像添加噪声"""
    def __init__(self, modes, p=0.5):
        """
        Params:
        modes: list or tuple of strings
            添加噪声的类型，如 'gaussian', 'localvar', 'poisson', 'salt', 'pepper', 's&p', 'speckle'
        p: float
            执行该操作的概率
        """
        self.modes = modes
        self.p = p

    def __call__(self, image):
        if random.uniform(0, 1) < self.p:  # 按概率执行该操作
            img_arr = np.array(image)
            for mode in self.modes:
                img_arr = skimage.util.random_noise(img_arr, mode)

            img_pil = Image.fromarray((img_arr * 255).astype(np.uint8))

            return img_pil
        else:
            return image

# 添加一波噪音
modes = ['gaussian'] *10
# modes = ['gaussian', 'pepper', 'speckle']
random_noise = RandomNoise(modes, p=1.)
noisy_image = random_noise(image)
noisy_image.show()

# 添加很多噪音
# modes = ['gaussian', 'pepper', 'speckle'] * 10