| import torch
|
| import math
|
| import numbers
|
| import random
|
| import numpy as np
|
|
|
| from PIL import Image, ImageOps
|
| from scipy.ndimage.filters import gaussian_filter
|
| from matplotlib.pyplot import imshow, imsave
|
| from scipy.ndimage.interpolation import map_coordinates
|
| import cv2
|
| from scipy import ndimage
|
| from torchvision import transforms
|
| import PIL, PIL.ImageOps, PIL.ImageEnhance, PIL.ImageDraw
|
| from torch import nn
|
| from scipy.ndimage.interpolation import zoom
|
|
|
|
|
| def to_multilabel(pre_mask, classes = 2):
|
| mask = np.zeros((pre_mask.shape[0], pre_mask.shape[1], classes))
|
| mask[pre_mask == 1] = [0, 1]
|
| mask[pre_mask == 2] = [1, 1]
|
| return mask
|
|
|
|
|
| class add_salt_pepper_noise():
|
| def __call__(self, sample):
|
| image = sample['image']
|
| X_imgs_copy = np.asarray(image).copy()
|
|
|
| salt_vs_pepper = 0.2
|
| amount = 0.004
|
|
|
| num_salt = np.ceil(amount * X_imgs_copy.size * salt_vs_pepper)
|
| num_pepper = np.ceil(amount * X_imgs_copy.size * (1.0 - salt_vs_pepper))
|
|
|
| seed = random.random()
|
| if seed > 0.75:
|
|
|
| coords = [np.random.randint(0, i - 1, int(num_salt)) for i in X_imgs_copy.shape]
|
| X_imgs_copy[coords[0], coords[1], :] = 1
|
| elif seed > 0.5:
|
|
|
| coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in X_imgs_copy.shape]
|
| X_imgs_copy[coords[0], coords[1], :] = 0
|
| sample['image'] = X_imgs_copy
|
| return sample
|
|
|
| class adjust_light():
|
| def __call__(self, sample):
|
| image = sample['image']
|
| seed = random.random()
|
| if seed > 0.5:
|
| gamma = random.random() * 3 + 0.5
|
| invGamma = 1.0 / gamma
|
| table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype(np.uint8)
|
| image = cv2.LUT(np.array(image).astype(np.uint8), table).astype(np.uint8)
|
| sample['image'] = image
|
| return sample
|
|
|
| class Brightness():
|
| def __init__(self, min_v, max_v):
|
| self.min_v = min_v
|
| self.max_v = max_v
|
|
|
| def __call__(self, img):
|
| v = self.min_v + float(self.max_v-self.min_v)*random.random()
|
| return PIL.ImageEnhance.Brightness(img).enhance(v)
|
|
|
| class Contrast():
|
| def __init__(self, min_v, max_v):
|
| self.min_v = min_v
|
| self.max_v = max_v
|
|
|
| def __call__(self, img):
|
| v = self.min_v + float(self.max_v-self.min_v)*random.random()
|
| return PIL.ImageEnhance.Contrast(img).enhance(v)
|
|
|
| class GaussianBlur(object):
|
| """blur a single image on CPU"""
|
| def __init__(self, kernel_size, num_channels):
|
| self.num_channels = num_channels
|
| radias = kernel_size // 2
|
| kernel_size = radias * 2 + 1
|
| self.blur_h = nn.Conv2d(num_channels, num_channels, kernel_size=(kernel_size, 1),
|
| stride=1, padding=0, bias=False, groups=num_channels)
|
| self.blur_v = nn.Conv2d(num_channels, num_channels, kernel_size=(1, kernel_size),
|
| stride=1, padding=0, bias=False, groups=num_channels)
|
| self.k = kernel_size
|
| self.r = radias
|
|
|
| self.blur = nn.Sequential(
|
| nn.ReflectionPad2d(radias),
|
| self.blur_h,
|
| self.blur_v
|
| )
|
|
|
| self.pil_to_tensor = transforms.ToTensor()
|
| self.tensor_to_pil = transforms.ToPILImage()
|
|
|
| def __call__(self, img):
|
| img = self.pil_to_tensor(img).unsqueeze(0)
|
|
|
| sigma = np.random.uniform(0.1, 2.0)
|
| x = np.arange(-self.r, self.r + 1)
|
| x = np.exp(-np.power(x, 2) / (2 * sigma * sigma))
|
| x = x / x.sum()
|
| x = torch.from_numpy(x).view(1, -1).repeat(self.num_channels, 1)
|
|
|
| self.blur_h.weight.data.copy_(x.view(self.num_channels, 1, self.k, 1))
|
| self.blur_v.weight.data.copy_(x.view(self.num_channels, 1, 1, self.k))
|
|
|
| with torch.no_grad():
|
| img = self.blur(img)
|
| img = img.squeeze()
|
|
|
| img = self.tensor_to_pil(img)
|
|
|
| return img
|
|
|
| class reverse_aug():
|
| """blur a single image on CPU"""
|
| def __init__(self, kernel_size, num_channels, min_v, max_v):
|
| self.min_v = min_v
|
| self.max_v = max_v
|
| self.num_channels = num_channels
|
| radias = kernel_size // 2
|
| kernel_size = radias * 2 + 1
|
| self.blur_h = nn.Conv2d(num_channels, num_channels, kernel_size=(kernel_size, 1),
|
| stride=1, padding=0, bias=False, groups=num_channels)
|
| self.blur_v = nn.Conv2d(num_channels, num_channels, kernel_size=(1, kernel_size),
|
| stride=1, padding=0, bias=False, groups=num_channels)
|
| self.k = kernel_size
|
| self.r = radias
|
|
|
| self.blur = nn.Sequential(
|
| nn.ReflectionPad2d(radias),
|
| self.blur_h,
|
| self.blur_v
|
| )
|
|
|
| self.pil_to_tensor = transforms.ToTensor()
|
| self.tensor_to_pil = transforms.ToPILImage()
|
|
|
| def __call__(self, img1, img2):
|
| v = self.min_v + float(self.max_v-self.min_v)*random.random()
|
| img1 = PIL.ImageEnhance.Brightness(img1).enhance(v)
|
| img2 = PIL.ImageEnhance.Brightness(img2).enhance(2-v)
|
| v = self.min_v + float(self.max_v-self.min_v)*random.random()
|
| img1 = PIL.ImageEnhance.Contrast(img1).enhance(v)
|
| img2 = PIL.ImageEnhance.Contrast(img2).enhance(2-v)
|
|
|
| img1 = self.pil_to_tensor(img1).unsqueeze(0)
|
| sigma = np.random.uniform(0.1, 2.0)
|
| x = np.arange(-self.r, self.r + 1)
|
| x = np.exp(-np.power(x, 2) / (2 * sigma * sigma))
|
| x = x / x.sum()
|
| x = torch.from_numpy(x).view(1, -1).repeat(self.num_channels, 1)
|
| self.blur_h.weight.data.copy_(x.view(self.num_channels, 1, self.k, 1))
|
| self.blur_v.weight.data.copy_(x.view(self.num_channels, 1, 1, self.k))
|
| with torch.no_grad():
|
| img1 = self.blur(img1)
|
| img1 = img1.squeeze()
|
| img1 = self.tensor_to_pil(img1)
|
|
|
| return img1, img2
|
|
|
|
|
| class eraser():
|
| def __call__(self, sample, s_l=0.02, s_h=0.06, r_1=0.3, r_2=0.6, v_l=0, v_h=255, pixel_level=False):
|
| image = sample['image']
|
| img_h, img_w, img_c = image.shape
|
|
|
|
|
| if random.random() > 0.5:
|
| return sample
|
|
|
| while True:
|
| s = np.random.uniform(s_l, s_h) * img_h * img_w
|
| r = np.random.uniform(r_1, r_2)
|
| w = int(np.sqrt(s / r))
|
| h = int(np.sqrt(s * r))
|
| left = np.random.randint(0, img_w)
|
| top = np.random.randint(0, img_h)
|
|
|
| if left + w <= img_w and top + h <= img_h:
|
| break
|
|
|
| if pixel_level:
|
| c = np.random.uniform(v_l, v_h, (h, w, img_c))
|
| else:
|
| c = np.random.uniform(v_l, v_h)
|
|
|
| image[top:top + h, left:left + w, :] = c
|
| sample['image'] = image
|
| return sample
|
|
|
| class elastic_transform():
|
| """Elastic deformation of images as described in [Simard2003]_.
|
| .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
|
| Convolutional Neural Networks applied to Visual Document Analysis", in
|
| Proc. of the International Conference on Document Analysis and
|
| Recognition, 2003.
|
| """
|
|
|
|
|
|
|
| def __call__(self, sample):
|
| image, label = sample['image'], sample['label']
|
| alpha = image.size[1] * 2
|
| sigma = image.size[1] * 0.08
|
| random_state = None
|
| seed = random.random()
|
| if seed > 0.5:
|
|
|
| assert len(image.size) == 2
|
|
|
| image_channel = len(np.array(image).shape)
|
| label_channel = len(np.array(label).shape)
|
|
|
| if random_state is None:
|
| random_state = np.random.RandomState(None)
|
|
|
| shape = image.size[0:2]
|
| dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
|
| dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
|
|
|
| x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
|
| indices = np.reshape(x + dx, (-1, 1)), np.reshape(y + dy, (-1, 1))
|
|
|
| if image_channel == 3:
|
| transformed_image = np.zeros([image.size[0], image.size[1], 3])
|
| for i in range(3):
|
|
|
| transformed_image[:, :, i] = map_coordinates(np.array(image)[:, :, i], indices, order=1).reshape(shape)
|
|
|
| elif image_channel == 2:
|
| transformed_image = np.zeros([image.size[0], image.size[1]])
|
| transformed_image[:, :] = map_coordinates(np.array(image)[:, :], indices, order=1).reshape(shape)
|
| if label is not None:
|
| if label_channel == 3:
|
| transformed_label = np.zeros([label.size[0], label.size[1], 3])
|
| for i in range(3):
|
| transformed_label[:, :, i] = map_coordinates(np.array(label)[:, :, i], indices, order=0, mode='nearest', prefilter=False).reshape(shape)
|
| elif label_channel == 2:
|
| transformed_label = np.zeros([label.size[0], label.size[1]])
|
| transformed_label[:, :] = map_coordinates(np.array(label)[:, :], indices, order=0, mode='nearest', prefilter=False).reshape(shape)
|
|
|
| else:
|
| transformed_label = None
|
| transformed_image = transformed_image.astype(np.uint8)
|
|
|
| if label is not None:
|
| transformed_label = transformed_label.astype(np.uint8)
|
| sample['image'] = Image.fromarray(transformed_image)
|
| sample['label'] = transformed_label
|
| return sample
|
|
|
| class cutout():
|
|
|
| def __init__(self):
|
| self.p=0.5
|
| self.size_min=0.02
|
| self.size_max=0.4
|
| self.ratio_1=0.3
|
| self.ratio_2=1/0.3
|
| self.value_min=0
|
| self.value_max=255
|
| self.pixel_level=True
|
|
|
| def __call__(self, sample):
|
| if random.random() < self.p:
|
| img, mask = sample['image'], sample['label']
|
| img = np.array(img)
|
| mask = np.array(mask)
|
|
|
| img_h, img_w = img.shape[0], img.shape[1]
|
| img_channel = len(img.shape)
|
|
|
| while True:
|
| size = np.random.uniform(self.size_min, self.size_max) * img_h * img_w
|
| ratio = np.random.uniform(self.ratio_1, self.ratio_2)
|
| erase_w = int(np.sqrt(size / ratio))
|
| erase_h = int(np.sqrt(size * ratio))
|
| x = np.random.randint(0, img_w)
|
| y = np.random.randint(0, img_h)
|
|
|
| if x + erase_w <= img_w and y + erase_h <= img_h:
|
| break
|
|
|
| if self.pixel_level:
|
| if img_channel == 3:
|
| value = np.random.uniform(self.value_min, self.value_max, (erase_h, erase_w, img.shape[2]))
|
| elif img_channel == 2:
|
| value = np.random.uniform(self.value_min, self.value_max, (erase_h, erase_w))
|
| else:
|
| value = np.random.uniform(self.value_min, self.value_max)
|
|
|
| img[y:y + erase_h, x:x + erase_w] = value
|
| mask[y:y + erase_h, x:x + erase_w] = 255
|
|
|
|
|
|
|
| sample['image'] = Image.fromarray(img.astype(np.uint8))
|
| sample['label'] = mask
|
|
|
| return sample
|
|
|
|
|
|
|
|
|
| class RandomCrop(object):
|
| def __init__(self, size, padding=0):
|
| if isinstance(size, numbers.Number):
|
| self.size = (int(size), int(size))
|
| else:
|
| self.size = size
|
| self.padding = padding
|
|
|
| def __call__(self, sample):
|
| img, mask = sample['image'], sample['label']
|
|
|
| w, h = img.size
|
| if self.padding > 0 or w < self.size[0] or h < self.size[1]:
|
| padding = np.maximum(self.padding,np.maximum((self.size[0]-w)//2+5,(self.size[1]-h)//2+5))
|
| img = ImageOps.expand(img, border=padding, fill=0)
|
| mask = ImageOps.expand(mask, border=padding, fill=255)
|
|
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
| w, h = img.size
|
| th, tw = self.size
|
| if w == tw and h == th:
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return sample
|
|
|
|
|
|
|
|
|
| x1 = random.randint(0, w - tw)
|
| y1 = random.randint(0, h - th)
|
| img = img.crop((x1, y1, x1 + tw, y1 + th))
|
| mask = mask.crop((x1, y1, x1 + tw, y1 + th))
|
|
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return sample
|
|
|
|
|
| class CenterCrop(object):
|
| def __init__(self, size):
|
| if isinstance(size, numbers.Number):
|
| self.size = (int(size), int(size))
|
| else:
|
| self.size = size
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
|
|
|
|
| w, h = img.size
|
| th, tw = self.size
|
| x1 = int(round((w - tw) / 2.))
|
|
|
| y1 = int(round((h - th) / 2.))
|
| img = img.crop((x1, y1, x1 + tw, y1 + th))
|
| mask = mask.crop((x1, y1, x1 + tw, y1 + th))
|
|
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': sample['img_name']}
|
|
|
|
|
| class RandomFlip(object):
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| if random.random() < 0.5:
|
| img = img.transpose(Image.FLIP_LEFT_RIGHT)
|
| mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
|
| if random.random() < 0.5:
|
| img = img.transpose(Image.FLIP_TOP_BOTTOM)
|
| mask = mask.transpose(Image.FLIP_TOP_BOTTOM)
|
|
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return sample
|
|
|
| class RandomHorizontalFlip(object):
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| if random.random() < 0.5:
|
| img = img.transpose(Image.FLIP_LEFT_RIGHT)
|
| mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
|
|
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return sample
|
|
|
|
|
| class FixedResize(object):
|
| def __init__(self, size):
|
| self.size = tuple(reversed(size))
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| name = sample['img_name']
|
|
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
| img = img.resize(self.size, Image.BILINEAR)
|
| mask = mask.resize(self.size, Image.NEAREST)
|
|
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': name}
|
|
|
|
|
| class Scale(object):
|
| def __init__(self, size):
|
| if isinstance(size, numbers.Number):
|
| self.size = (int(size), int(size))
|
| else:
|
| self.size = size
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
| w, h = img.size
|
|
|
| if (w >= h and w == self.size[1]) or (h >= w and h == self.size[0]):
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': sample['img_name']}
|
| oh, ow = self.size
|
| img = img.resize((ow, oh), Image.BILINEAR)
|
| mask = mask.resize((ow, oh), Image.NEAREST)
|
|
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': sample['img_name']}
|
|
|
|
|
| class RandomSizedCrop(object):
|
| def __init__(self, size):
|
| self.size = size
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| name = sample['img_name']
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
| for attempt in range(10):
|
| area = img.size[0] * img.size[1]
|
| target_area = random.uniform(0.45, 1.0) * area
|
| aspect_ratio = random.uniform(0.5, 2)
|
|
|
| w = int(round(math.sqrt(target_area * aspect_ratio)))
|
| h = int(round(math.sqrt(target_area / aspect_ratio)))
|
|
|
| if random.random() < 0.5:
|
| w, h = h, w
|
|
|
| if w <= img.size[0] and h <= img.size[1]:
|
| x1 = random.randint(0, img.size[0] - w)
|
| y1 = random.randint(0, img.size[1] - h)
|
|
|
| img = img.crop((x1, y1, x1 + w, y1 + h))
|
| mask = mask.crop((x1, y1, x1 + w, y1 + h))
|
| assert (img.size == (w, h))
|
|
|
| img = img.resize((self.size, self.size), Image.BILINEAR)
|
| mask = mask.resize((self.size, self.size), Image.NEAREST)
|
|
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': name}
|
|
|
|
|
| scale = Scale(self.size)
|
| crop = CenterCrop(self.size)
|
| sample = crop(scale(sample))
|
| return sample
|
|
|
|
|
| class RandomRotate(object):
|
| def __init__(self, size=512):
|
| self.degree = random.randint(1, 4) * 90
|
| self.size = size
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
|
|
| seed = random.random()
|
| if seed > 0.5:
|
| rotate_degree = self.degree
|
| img = img.rotate(rotate_degree, Image.BILINEAR, expand=0)
|
| mask = mask.rotate(rotate_degree, Image.NEAREST, expand=255)
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return sample
|
|
|
| class RandomScaleRotate(object):
|
| def __init__(self, size=512, left=-20, right=20, fillcolor=255):
|
| self.size = size
|
| self.left = left
|
| self.right = right
|
| self.fillcolor = fillcolor
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
|
|
| seed = random.random()
|
| if seed > 0.5:
|
| rotate_degree = random.randint(self.left, self.right)
|
| img = img.rotate(rotate_degree, Image.BILINEAR)
|
| mask = mask.rotate(rotate_degree, Image.NEAREST, fillcolor=self.fillcolor)
|
|
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return sample
|
|
|
|
|
| class RandomScaleCrop(object):
|
| def __init__(self, size):
|
| self.size = size
|
|
|
| self.crop = RandomCrop(self.size)
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
|
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
|
|
| seed = random.random()
|
| if seed > 0.5:
|
| w = int(random.uniform(1, 1.5) * img.size[0])
|
| h = int(random.uniform(1, 1.5) * img.size[1])
|
|
|
| img, mask = img.resize((w, h), Image.BILINEAR), mask.resize((w, h), Image.NEAREST)
|
| sample['image'] = img
|
| sample['label'] = mask
|
| return self.crop(sample)
|
|
|
|
|
| class ResizeImg(object):
|
| def __init__(self, size):
|
| self.size = size
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| name = sample['img_name']
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
|
|
| img = img.resize((self.size, self.size))
|
|
|
|
|
| sample = {'image': img, 'label': mask, 'img_name': name}
|
| return sample
|
|
|
|
|
| class Resize(object):
|
| def __init__(self, size):
|
| self.size = size
|
|
|
| def __call__(self, sample):
|
| img = sample['image']
|
| mask = sample['label']
|
| name = sample['img_name']
|
| assert img.width == mask.width
|
| assert img.height == mask.height
|
|
|
| img = img.resize((self.size, self.size))
|
| mask = mask.resize((self.size, self.size))
|
|
|
| sample = {'image': img, 'label': mask, 'img_name': name}
|
| return sample
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| class Normalize(object):
|
| """Normalize a tensor image with mean and standard deviation.
|
| Args:
|
| mean (tuple): means for each channel.
|
| std (tuple): standard deviations for each channel.
|
| """
|
| def __init__(self, mean=(0., 0., 0.), std=(1., 1., 1.)):
|
| self.mean = mean
|
| self.std = std
|
|
|
| def __call__(self, sample):
|
| img = np.array(sample['image']).astype(np.float32)
|
| mask = np.array(sample['label']).astype(np.float32)
|
| img /= 255.0
|
| img -= self.mean
|
| img /= self.std
|
|
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': sample['img_name']}
|
|
|
|
|
| class GetBoundary(object):
|
| def __init__(self, width = 5):
|
| self.width = width
|
| def __call__(self, mask):
|
| cup = mask[:, :, 0]
|
| disc = mask[:, :, 1]
|
| dila_cup = ndimage.binary_dilation(cup, iterations=self.width).astype(cup.dtype)
|
| eros_cup = ndimage.binary_erosion(cup, iterations=self.width).astype(cup.dtype)
|
| dila_disc= ndimage.binary_dilation(disc, iterations=self.width).astype(disc.dtype)
|
| eros_disc= ndimage.binary_erosion(disc, iterations=self.width).astype(disc.dtype)
|
| cup = dila_cup + eros_cup
|
| disc = dila_disc + eros_disc
|
| cup[cup==2]=0
|
| disc[disc==2]=0
|
| size = mask.shape
|
|
|
| boundary = (cup + disc) > 0
|
| return boundary.astype(np.uint8)
|
|
|
|
|
| class Normalize_tf(object):
|
| """Normalize a tensor image with mean and standard deviation.
|
| Args:
|
| mean (tuple): means for each channel.
|
| std (tuple): standard deviations for each channel.
|
| """
|
| def __init__(self, mean=(0., 0., 0.), std=(1., 1., 1.), dataRange = [-1, 1]):
|
| self.mean = mean
|
| self.std = std
|
| self.get_boundary = GetBoundary()
|
| self.dataRange = dataRange
|
|
|
| def __call__(self, sample):
|
| img = np.array(sample['image']).astype(np.float32)
|
|
|
| if self.dataRange == [-1, 1]:
|
| img /= 127.5
|
| img -= 1.0
|
| elif self.dataRange == [0, 1]:
|
| img /= 255
|
| else:
|
| assert(0)
|
| if 'strong_aug' in sample.keys():
|
| strong = np.array(sample['strong_aug']).astype(np.float32)
|
| if self.dataRange == [-1, 1]:
|
| strong /= 127.5
|
| strong -= 1.0
|
| elif self.dataRange == [0, 1]:
|
| strong /= 255
|
| else:
|
| assert(0)
|
| sample['strong_aug'] = strong
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| sample['image'] = img
|
|
|
| return sample
|
|
|
|
|
| class Normalize_cityscapes(object):
|
| """Normalize a tensor image with mean and standard deviation.
|
| Args:
|
| mean (tuple): means for each channel.
|
| std (tuple): standard deviations for each channel.
|
| """
|
| def __init__(self, mean=(0., 0., 0.)):
|
| self.mean = mean
|
|
|
| def __call__(self, sample):
|
| img = np.array(sample['image']).astype(np.float32)
|
| mask = np.array(sample['label']).astype(np.float32)
|
| img -= self.mean
|
| img /= 255.0
|
|
|
| return {'image': img,
|
| 'label': mask,
|
| 'img_name': sample['img_name']}
|
|
|
| def ToMultiLabel(dc):
|
| new_dc = np.zeros([3])
|
| for i in range(new_dc.shape[0]):
|
| if i == dc:
|
| new_dc[i] = 1
|
| return new_dc
|
|
|
| def SoftLable(label):
|
| new_label = label.copy()
|
| label = list(label)
|
| index = label.index(1)
|
| new_label[index] = 0.8+random.random()*0.2
|
| accelarate = new_label[index]
|
| for i in range(len(label)):
|
| if i != index:
|
| if i == len(label) - 1:
|
| new_label[i] = 1 - accelarate
|
| else:
|
| new_label[i] = random.random()*(1-accelarate)
|
| accelarate += new_label[i]
|
| return new_label
|
|
|
| class ToTensor(object):
|
| """Convert ndarrays in sample to Tensors."""
|
| def __init__(self, low_res=0, unet_size=0):
|
| self.low_res = low_res
|
| self.unet_size = unet_size
|
|
|
| def __call__(self, sample):
|
|
|
|
|
|
|
| if len(np.array(sample['image']).shape) == 2:
|
| sample['image'] = np.expand_dims(np.array(sample['image']).astype(np.float32), 2)
|
|
|
|
|
| img = np.array(sample['image']).astype(np.float32).transpose((2, 0, 1))
|
| map = np.array(sample['label']).astype(np.uint8)
|
| img_x, img_y = img.shape[1:]
|
| x, y = map.shape
|
| if self.low_res > 0:
|
| low_res_label = zoom(map, (self.low_res / x, self.low_res / y), order=0)
|
| low_res_label = torch.from_numpy(low_res_label).float()
|
| sample['low_res_label']=low_res_label
|
| if self.unet_size > 0:
|
| unet_size_img = zoom(img, (1, self.unet_size / img_x, self.unet_size / img_y), order=0)
|
| unet_size_img = torch.from_numpy(unet_size_img).float()
|
| sample['unet_size_img']=unet_size_img
|
| unet_size_label = zoom(map, (self.unet_size / x, self.unet_size / y), order=0)
|
| unet_size_label = torch.from_numpy(unet_size_label).float()
|
| sample['unet_size_label']=unet_size_label
|
| if 'strong_aug' in sample.keys():
|
| if len(np.array(sample['strong_aug']).shape) == 2:
|
| sample['strong_aug'] = np.expand_dims(np.array(sample['strong_aug']).astype(np.float32), 2)
|
| strong = np.array(sample['strong_aug']).astype(np.float32).transpose((2, 0, 1))
|
| if self.unet_size > 0:
|
| unet_size_strong_aug = zoom(strong, (1, self.unet_size / img_x, self.unet_size / img_y), order=0)
|
| unet_size_strong_aug = torch.from_numpy(unet_size_strong_aug).float()
|
| sample['unet_size_strong_aug']=unet_size_strong_aug
|
| strong = torch.from_numpy(strong).float()
|
| sample['strong_aug'] = strong
|
| img = torch.from_numpy(img).float()
|
| map = torch.from_numpy(map).float()
|
| sample['image']=img
|
| sample['label']=map
|
|
|
|
|
| return sample |
