数据增强
- 随机旋转
- 随机缩放
- 随机偏移
- 随机扭曲
- 随机噪声
- 随机亮度对比度调整
模板代码Data.py
python
import cv2
import numpy as np
import torch.utils.data as data
from PIL import Image
from PIL import ImageOps
from PIL import ImageEnhance
Image.LOAD_TRUNCATED_IMAGES = True # 防止图片过大，加载不完整
def preprocess_input(x):
x /= 127.5
x -= 1.
return x
def random_distort(image, max_distort=0.3):
"""
Apply random distortion to the image.
:param image: PIL Image
:param max_distort: Maximum distortion factor
:return: Distorted image as a PIL Image
"""
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
w, h = image.size
distort_x = np.random.uniform(-max_distort, max_distort) * w
distort_y = np.random.uniform(-max_distort, max_distort) * h
src_points = np.float32([[0, 0], [w, 0], [0, h], [w, h]])
dst_points = np.float32([[distort_x, distort_y], [w - distort_x, distort_y], [distort_x, h - distort_y], [w - distort_x, h - distort_y]])
M = cv2.getPerspectiveTransform(src_points, dst_points)
distorted_image = cv2.warpPerspective(np.array(image), M, (w, h), borderValue=[128, 128, 128])
return Image.fromarray(distorted_image)
def cvtColor(image):
if len(np.shape(image)) == 3 and np.shape(image)[-2] == 3:
return image
else:
image = image.convert('RGB')
return image
class DataGenerator(data.Dataset):
def __init__(self, annotation_lines, inpt_shape, random=True, window_size=5):
self.window_size = window_size
self.targets = [int(line.split(';')[0]) for line in annotation_lines] # 收集所有标签
self.annotation_lines = annotation_lines
self.input_shape = inpt_shape
self.random = random
def __len__(self):
return len(self.annotation_lines)
def __getitem__(self, index):
annotation_path = self.annotation_lines[index].split(';')[1].split()[0]
image = Image.open(annotation_path)
image = self.get_random_data(image, self.input_shape, index, random=self.random)
image = np.transpose(preprocess_input(np.array(image).astype(np.float32)), [2, 0, 1])
y = int(self.annotation_lines[index].split(';')[0])
return image, y, annotation_path
def rand(self, a=0, b=1):
return np.random.rand() * (b - a) + a
def get_random_data(self, image, inpt_shape, index, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True):
image = cvtColor(image)
iw, ih = image.size
h, w = inpt_shape
if not random:
scale = min(w / iw, h / ih) # 缩放比例
nw = int(iw * scale) # 新的宽度
nh = int(ih * scale) # 新的高度
dx = (w - nw) // 2 # x方向的偏移量
dy = (h - nh) // 2 # y方向的偏移量
image = image.resize((nw, nh), Image.BICUBIC)
new_image = Image.new('RGB', (w, h), (128, 128, 128))
new_image.paste(image, (dx, dy))
image_data = np.array(new_image, np.float32)
return image_data
new_ar = w / h * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter) # 随机长宽比
scale = self.rand(.75, 1.25) # 随机缩放
if new_ar < 1:
nh = int(scale * h)
nw = int(nh * new_ar)
else:
nw = int(scale * w)
nh = int(nw / new_ar)
image = image.resize((nw, nh), Image.BICUBIC)
#将图像多余的部分加上灰条
# dx=int(self.rand(0,w-nw)) #x方向的偏移量
# dy=int(self.rand(0,h-nh)) #y方向的偏移量
dx = (w - nw) // 2 # x方向的偏移量
dy = (h - nh) // 2 # y方向的偏移量
new_image=Image.new('RGB',(w,h),(128,128,128))
new_image.paste(image,(dx,dy))
image=new_image
# # 随机缩小
# shrink_factor_w = self.rand(0.5, 1) # 宽度缩小因子
# shrink_factor_h = self.rand(0.5, 1) # 高度缩小因子
# # 计算新尺寸
# nw, nh = int(iw * shrink_factor_w), int(ih * shrink_factor_h)
# image = image.resize((nw, nh), Image.BICUBIC)
# # 计算填充量
# pad_w = (w - nw) // 2
# pad_h = (h - nh) // 2
# new_image = Image.new('RGB', (w, h), (128, 128, 128))
# new_image.paste(image, (pad_w, pad_h))
# image = new_image
# # 随机噪声
# if self.rand() < 0.5:
# image_np = np.array(image) # 将图像转换为numpy数组
# noise = np.random.normal(0, 0.02, (h, w, 3))
# image_np = image_np + noise
# image = Image.fromarray(image_np.astype('uint8')) # 将图像转换回PIL Image对象
#
# # 随机亮度、对比度调整
# if self.rand() < 1:
# brightness_factor = self.rand(0.8, 1.2)
# image = ImageEnhance.Brightness(image).enhance(brightness_factor)
# if self.rand() < 1:
# contrast_factor = self.rand(0.8, 1.2)
# image = ImageEnhance.Contrast(image).enhance(contrast_factor)
# 垂直翻转
if self.rand() < .5:
image = ImageOps.flip(image)
# 翻转图像
flip = self.rand() < .5
if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)
rotate = self.rand() < .5
if rotate:
angle = np.random.randint(-30, 30)
a, b = w / 2, h / 2
M = cv2.getRotationMatrix2D((a, b), angle, 1)
image = cv2.warpAffine(np.array(image), M, (w, h), borderValue=[128, 128, 128]) # 旋转图像
if self.rand() < .5:
image = random_distort(image, max_distort=0.3) # 随机扭曲图像
image_data = np.array(image, np.float32)
# # 色域扭曲
# hue = self.rand(-hue, hue)
# sat = self.rand(1, sat) if self.rand() < .5 else 1 / self.rand(1, sat)
# val = self.rand(1, val) if self.rand() < .5 else 1 / self.rand(1, val)
# x = cv2.cvtColor(np.array(image, np.float32) / 255, cv2.COLOR_RGB2HSV) # 颜色空间转换
# x[..., 1] *= sat
# x[..., 2] *= val
# x[x[:, :, 0] > 360, 0] = 360
# x[:, :, 1:][x[:, :, 1:] > 1] = 1
# x[x < 0] = 0
# image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB) * 255
return image_data