Pytorch对图像增强的python实现

直接上代码
介绍的也比较详细

import torchvision.transforms.functional as TF
from PIL import Image
import os
from skimage.util import random_noise
from skimage.filters import gaussian
import cv2

# 擦除图像中某块区域
def erase_and_save(image_path, target_dir, position, size):
    """
    按照指定的位置和长宽擦除
    :param image_path: 输入图像路径
    :param target_dir: 目标图像路径
    :param position: 擦除的左上角坐标
    :param size: 擦除的长宽值
    :return: 返回擦除后的图像
    """
    image = TF.to_tensor(Image.open(image_path))
    erased_image = TF.to_pil_image(TF.erase(img=image,
                            i=position[0],
                            j=position[1],
                            h=size[0],
                            w=size[1],
                            v=1))
    erased_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4]))+"_erased_.jpg")
    return erased_image

# 任意角度旋转
def rotate_and_save(image_path, target_dir, angles_list):
    """
    按照一系列的指定角度旋转输入图像并保存
    :param image_path:输入图像路径
    :param target_dir:目标图片存储目录
    :param angles_list: 需要对输入图像进行多少角度的旋转
    :return: 完成后以列表形式返回所有的旋转后图像
    """
    image = Image.open(image_path)

    # 创建旋转后的图像容器
    image_list = []
    for angle in angles_list:
        rotated_image = TF.rotate(img=image, angle=angle, resample=Image.NEAREST)
        # 按照原始image_path的图片名称作为目标图像名称
        # 获取图片名称并去掉后缀".jpg"
        rotated_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4])) + "_" + str(angle) + "_.jpg")
        image_list.append(rotated_image)
    return image_list

# 水平和垂直翻转
def vflip_and_save(image_path, target_dir):
    """
    垂直旋转翻转图像，并保存
    :param image_path: 输入图像路径
    :param target_path: 目标图片存储目录，如"d/"
    :return: 返回旋转后的图片
    """
    image = Image.open(image_path)
    vertical_image = TF.vflip(img=image)
    vertical_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4]))+"_v_.jpg")

    return vertical_image

def hflip_and_save(image_path, target_dir):
    """
    垂直旋转翻转图像，并保存
    :param image_path: 输入图像路径
    :param target_path: 目标图片存储目录，如"d/"
    :return: 返回旋转后的图片
    """
    image = Image.open(image_path)
    horizontal_image = TF.hflip(img=image)
    horizontal_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4]))+"_h_.jpg")

    return horizontal_image


def brightness_and_save(image_path, target_dir):
    
    '''
    图像变暗变亮
    :param image_path: 输入图像路径
    :param target_dir: 目标图片存储目录，如"d/"
    :return: 变换后的图片
    '''
    # brightness_factor=[0,2]  ==1 原图相不发生变化， 小于1变亮，大于1变暗。
    image = Image.open(image_path)
    brightness_image = TF.adjust_brightness(image, brightness_factor=0.5)
    brightness_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4]))+"_brightness_.jpg")

    return brightness_image


def adjust_hue_and_save(image_path, target_dir):
    
    image = Image.open(image_path)
    adjust_hue_image = TF.adjust_hue(image, hue_factor=0.5)
    adjust_hue_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4])) + "_adjust_.jpg")

    return adjust_hue_image


# 非数据增强方法，而是图像增强
def gamma_and_save(image_path, target_dir, gamma_value):
    """
    进行伽马变换并保存
    :param image_path: 输入图片路径
    :param target_dir: 目标图像存储目录
    :param gamma_value: 伽马值
    :return: 伽马变换后的图像
    """
    image = Image.open(image_path)
    gamma_image = TF.adjust_gamma(img=image, gamma=gamma_value)
    gamma_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4]))+"_gamma_.jpg")
    return gamma_image


if __name__ == '__main__':
    # 测试效果，若批量图像可以写个循环。数据太多的话用多线程处理
    image_path = '/home/bitwenzi/1/00000001.jpg'
    target_dir = '/home/bitwenzi/2'
    rotate_and_save(image_path, target_dir, angles_list=[10,20,30,40])
    vflip_and_save(image_path, target_dir)
    gamma_and_save(image_path, target_dir, gamma_value=0.1)
    adjust_hue_and_save(image_path, target_dir)
    brightness_and_save(image_path, target_dir)
    hflip_and_save(image_path, target_dir)
    # res1 = erase_and_save(image_path, target_dir, position=[200,200], size=[50,50])