File size: 8,225 Bytes
edcf5ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 |
"""
dual augmentation on both images and their masks Script ver： Apr 10th 11:20
"""
import random
import numpy as np
import cv2
from PIL import Image
from torchvision import transforms
from utils.tools import to_2tuple
class DualCompose: # fit pytorch transform
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, mask=None):
# process the cv2 transformation first
for t in self.transforms:
image, mask = t(image, mask)
# NOTICE 转回图片 值总和还变成了cv2 numpy的1/255
# Then, Transform cv2 BGR image to PIL RGB image
# BGR -> RGB channel
b, g, r = cv2.split(image)
image = cv2.merge([r, g, b])
b, g, r = cv2.split(mask)
mask = cv2.merge([r, g, b])
# Image.fromarray make the 0-255 to PIL 0-1 range values
return Image.fromarray(np.uint8(image)), Image.fromarray(np.uint8(mask))
class DualImageTransform:
# Transform cv2 BGR image to PIL RGB image
def __init__(self):
pass
def __call__(self, image, mask=None):
# BGR -> RGB channel
b, g, r = cv2.split(image)
image = cv2.merge([r, g, b])
b, g, r = cv2.split(mask)
mask = cv2.merge([r, g, b])
# Image.fromarray make the 0-255 to PIL 0-1 range values
return Image.fromarray(np.uint8(image)), Image.fromarray(np.uint8(mask))
class Dual_RandomHorizontalFlip:
"""
Random horizontal flip.
image shape: (height, width, channels)
mask shape: (height, width)
possibility: possibility for flip
"""
def __init__(self, possibility=0.5):
assert isinstance(possibility, (int, float))
self.possibility = possibility
def __call__(self, image, mask):
if random.random() <= self.possibility:
image = np.flip(image, axis=1)
mask = np.flip(mask, axis=1)
return image, mask
class Dual_RandomVerticalFlip:
"""
Random vertical flip.
image shape: (height, width, channels)
mask shape: (height, width)
possibility: possibility for flip
"""
def __init__(self, possibility=0.5):
assert isinstance(possibility, (int, float))
self.possibility = possibility
def __call__(self, image, mask):
if random.random() <= self.possibility:
image = np.flip(image, axis=0)
mask = np.flip(mask, axis=0)
return image, mask
class Dual_Rotate:
"""
Random rotation.
image shape: (height, width, channels)
mask shape: (height, width)
possibility: possibility for rotate
range: range of rotation angles
"""
def __init__(self, possibility=0.5, range=20):
self.possibility = possibility
self.range = range
def __call__(self, image, mask):
# 这里cv2读到的是反的，因此这里是height, width而不是width，height，图片input不是正方形时会有严重后果
height, width = image.shape[:2]
if random.random() <= self.possibility:
angle = np.random.randint(0, self.range)
center = (width // 2, height // 2)
# 得到旋转矩阵，第一个参数为旋转中心，第二个参数为旋转角度，第三个参数为旋转之前原图像缩放比例
M = cv2.getRotationMatrix2D(center, -angle, 1)
# 进行仿射变换，第一个参数图像，第二个参数是旋转矩阵，第三个参数是变换之后的图像大小
image = cv2.warpAffine(image, M, (width, height))
mask = cv2.warpAffine(mask.astype(np.uint8), M, (width, height))
return image.astype(np.uint8), mask.astype(np.int)
def Four_step_dual_augmentation(data_augmentation_mode=0, edge_size=384):
"""
Get data augmentation methods
Dual_transform : Transform CV2 images and their mask by Rotate, RandomHorizontalFlip, etc.
DualImage : Transform CV2 images and their mask to PIL images
train_domain_transform : transforms.ColorJitter on PIL images
transform: PIL crop, resize and to Tensor
USAGE:
IN Train:
image, mask = self.Dual_transform(image, mask)
# image color jitter shifting
image = self.train_domain_transform(image)
# crop + resize
image = self.transform(image)
IN Val $ Test:
# 0/255 mask -> binary mask
image, mask = self.DualImage(image, mask)
# crop + resize
image = self.transform(image)
"""
edge_size = to_2tuple(edge_size)
if data_augmentation_mode == 0: # ROSE + MARS
# apply the on-time synchornized transform on image and mask togather
Dual_transform = DualCompose([
Dual_Rotate(possibility=0.8, range=180),
Dual_RandomHorizontalFlip(),
Dual_RandomVerticalFlip(),
])
# val & test use DualImage to convert PIL Image
DualImage = DualImageTransform()
# ColorJitter for image only
train_domain_transform = transforms.Compose([
# HSL shift operation
transforms.ColorJitter(brightness=0.15, contrast=0.3, saturation=0.3, hue=0.06),
])
# lastly, the synchornized separate transform
transform = transforms.Compose([
transforms.CenterCrop(700), # center area for classification
transforms.Resize(edge_size),
transforms.ToTensor(), # hwc -> chw tensor
])
elif data_augmentation_mode == 1: # Cervical
# apply the on-time synchornized transform on image and mask togather
Dual_transform = DualCompose([
Dual_Rotate(possibility=0.8, range=180),
Dual_RandomHorizontalFlip(),
Dual_RandomVerticalFlip(),
])
# val & test use DualImage to convert PIL Image
DualImage = DualImageTransform()
# ColorJitter for image only
train_domain_transform = transforms.Compose([
# HSL shift operation
transforms.ColorJitter(brightness=0.15, contrast=0.3, saturation=0.3, hue=0.06),
])
# lastly, the synchornized separate transform
transform = transforms.Compose([
transforms.Resize(edge_size),
transforms.ToTensor(), # hwc -> chw tensor
])
elif data_augmentation_mode == 2: #
# apply the on-time synchornized transform on image and mask togather
Dual_transform = DualCompose([
Dual_Rotate(possibility=0.8, range=180),
Dual_RandomHorizontalFlip(),
Dual_RandomVerticalFlip(),
])
# val & test use DualImage to convert PIL Image
DualImage = DualImageTransform()
# ColorJitter for image only
train_domain_transform = transforms.Compose([
# HSL shift operation
transforms.ColorJitter(brightness=0.15, contrast=0.3, saturation=0.3, hue=0.06),
])
# lastly, the synchornized separate transform
transform = transforms.Compose([
transforms.CenterCrop(360), # center area for classification
transforms.Resize(edge_size),
transforms.ToTensor(), # hwc -> chw tensor
])
elif data_augmentation_mode == 3: # for the squre input: just resize
# apply the on-time synchornized transform on image and mask togather
Dual_transform = DualCompose([
# Dual_Rotate(possibility=0.8, range=180),
Dual_RandomHorizontalFlip(),
Dual_RandomVerticalFlip(),
])
# val & test use DualImage to convert PIL Image
DualImage = DualImageTransform()
# ColorJitter for image only
train_domain_transform = transforms.Compose([
# HSL shift operation
transforms.ColorJitter(brightness=0.15, contrast=0.3, saturation=0.3, hue=0.06),
])
# lastly, the synchornized separate transform
transform = transforms.Compose([
transforms.Resize(edge_size),
transforms.ToTensor(), # hwc -> chw tensor
])
else:
print('no legal data augmentation is selected')
return -1
return Dual_transform, DualImage, train_domain_transform, transform
|