一、简介
不需要修改标签的数据增强有变明，变暗，hsv增强，color增强，cutout，模拟太阳光，雨水，雾等。
二、 代码
import random
import numpy as np
import os.path
import cv2
from PIL import Image
from PIL import ImageEnhance
from torchvision import transforms
import torch
orgin_dir = "E:/HPA-MTL-Ofilm/hpa-mtl-ofilm_debug/data/data_750/train/images"
aug_dir = "E:/HPA-MTL-Ofilm/hpa-mtl-ofilm_debug/data/data_750/train/aug_images/"
if not os.path.exists(aug_dir):
os.mkdir(aug_dir)
def bright_aug(image, percetage):
image_copy = image.copy()
w = image.shape[1]
h = image.shape[0]
#get brighter
for xi in range(0,w):
for xj in range(0,h):
image_copy[xj,xi,0] = np.clip(int(image[xj,xi,0]*percetage),a_max=255,a_min=0)
image_copy[xj,xi,1] = np.clip(int(image[xj,xi,1]*percetage),a_max=255,a_min=0)
image_copy[xj,xi,2] = np.clip(int(image[xj,xi,2]*percetage),a_max=255,a_min=0)
return image_copy
def dark_aug(image,percetage):
image_copy = image.copy()
w = image.shape[1]
h = image.shape[0]
#get darker
for xi in range(0,w):
for xj in range(0,h):
image_copy[xj,xi,0] = int(image[xj,xi,0]*percetage)
image_copy[xj,xi,1] = int(image[xj,xi,1]*percetage)
image_copy[xj,xi,2] = int(image[xj,xi,2]*percetage)
return image_copy
def blur_aug(image,):
#image
image_copy = cv2.GaussianBlur(image, (7, 7), 1.5) # cv2.GaussianBlur(图像，卷积核，标准差）
return image_copy
def hsv_aug(image, hgain=5, sgain=50, vgain=50):
# BGR->HSV
img_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.int16)
hsv_augs = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] # uniform 上界：1 下届：-1 尺寸：3
hsv_augs *= np.random.randint(0, 2, 3) # randint 上界：2 下届：0 尺寸：3
hsv_augs = hsv_augs.astype(np.int16) # int16
img_hsv[..., 0] = (img_hsv[..., 0] + hsv_augs[0]) % 180
img_hsv[..., 1] = np.clip(img_hsv[..., 1] + hsv_augs[1], 0, 255)
img_hsv[..., 2] = np.clip(img_hsv[..., 2] + hsv_augs[2], 0, 255)
# HSV->BGR
img_bgr = cv2.cvtColor(img_hsv.astype(image.dtype), cv2.COLOR_HSV2BGR)
return img_bgr
def random_color(image):
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) # 格式转换
random_factor = np.random.randint(0, 31) / 10 # 随机因子
color_image = ImageEnhance.Color(image).enhance(random_factor) # 饱和度
random_factor = np.random.randint(8, 15) / 10 # 随机因子
brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor) # 亮度
random_factor = np.random.randint(10, 21) / 10 # 随机因子
contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor) # 对比度
random_factor = np.random.randint(0, 31) / 10 # 随机因子
shape_image = ImageEnhance.Sharpness(contrast_image).enhance(random_factor) # 锐度
image_out = cv2.cvtColor(np.asarray(shape_image), cv2.COLOR_RGB2BGR) # 格式转换
return image_out
def sunlight_aug(img):
rows, cols = img.shape[:2]
centerX = np.random.randint(0, rows)
centerY = np.random.randint(0, cols)
radius = np.random.randint(50, min(rows / 2, cols / 2))
arr_rows = np.arange(rows).reshape(rows, 1)
arr_cols = np.arange(cols).reshape(1, cols)
arr = 1 - np.sqrt(((arr_cols - centerY) ** 2) + ((arr_rows - centerX) ** 2)) / radius
value_B = np.random.randint(255, 256)
value_G = np.random.randint(255, 256)
value_R = np.random.randint(255, 256)
arr_B = np.maximum(np.int16(value_B * arr), 0)
arr_G = np.maximum(np.int16(value_G * arr), 0)
arr_R = np.maximum(np.int16(value_R * arr), 0)
img = np.int16(img)
img[:, :, 0] += np.int16(arr_B / 10 * 8)
img[:, :, 1] += np.int16(arr_G / 10 * 8)
img[:, :, 2] += np.int16(arr_R / 10 * 8)
img = np.maximum(img, 0)
img = np.minimum(img, 255)
img = np.uint8(img)
return img
def cutout_aug(image, n_holes=25, height=320, width=320):
# 根据图像尺寸，调整小黑块大小
length_ratio = 0.2 * (height // 640 + width // 640)
length = int(35 * length_ratio)
# array -> tensor
image = transforms.ToTensor()(image)
h = image.size(1)
w = image.size(2)
mask = np.ones((h, w), np.float32)
for n in range(n_holes):
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - length // 2, 0, h)
y2 = np.clip(y + length // 2, 0, h)
x1 = np.clip(x - length // 2, 0, w)
x2 = np.clip(x + length // 2, 0, w)
mask[y1:y2, x1:x2] = 0
mask = torch.from_numpy(mask)
mask = mask.expand_as(image)
image = image * mask
# tensor -> array
image = image.mul(255).byte()
image = image.numpy().transpose((1, 2, 0))
return image.copy()
bright_aug_p = 0.3
dark_aug_p = 0.3
blur_aug_p = 0.3
hsv_p = 0.3
random_color_p = 0.3
sunlight_p = 0.3
cutout_p = 0.4
num_files = int(len(os.listdir(orgin_dir)))
num=0
for img_name in os.listdir(orgin_dir):
if img_name[len(img_name) - 4: len(img_name)] == '.png' or img_name[len(img_name) - 4: len(img_name)] == '.jpg' or img_name[len(img_name) - 4: len(img_name)] == '.bmp':
num+=1
# ----- read -----#
img_path = orgin_dir + '/' + img_name
img = cv2.imread(img_path)
# ----- bright -----#
if random.random() < bright_aug_p:
img=bright_aug(img, percetage=1.5)
# ----- dark -----#
if random.random() < dark_aug_p:
img=dark_aug(img, percetage=0.9)
# ----- blur -----#
if random.random() < blur_aug_p:
img=blur_aug(img)
# ----- hsv -----#
if random.random() < hsv_p:
img=hsv_aug(img)
# ----- color -----#
if random.random() < random_color_p:
img=random_color(img)
# ----- sunlight -----#
if random.random() < sunlight_p:
img=sunlight_aug(img)
# ----- cutput -----#
if random.random() < cutout_p:
img=cutout_aug(img)
# ----- save -----#
save_img_path = aug_dir + img_name[0:-4] + '.jpg'
cv2.imwrite(save_img_path, img)
print("已完成%d张图片增强，进度为%f %%" % (num, num * 100 / num_files))