| import torch, random | |
| import torchvision.transforms.functional as F | |
| import torchvision.transforms as tf | |
| import numpy as np | |
| from PIL import Image | |
| from typing import Tuple, List, Callable | |
| class Compose: | |
| def __init__(self, | |
| transforms: List[Callable], | |
| student_augs: bool = False): | |
| self.transforms = transforms | |
| self.student_augs = student_augs | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[torch.Tensor, torch.Tensor]: | |
| for transform in self.transforms: | |
| if pseudo is None: | |
| img, gt = transform(img, gt) | |
| else: | |
| img, gt, pseudo = transform(img, gt, pseudo) | |
| if self.student_augs: | |
| aimg = img.clone() | |
| aimg, _ = RandGaussianBlur()(aimg, gt) | |
| if 0.5 > random.random(): | |
| aimg, _ = ColorJitter()(aimg, gt) | |
| else: | |
| aimg, _ = MaskGrayscale()(aimg, gt) | |
| if pseudo is None and not self.student_augs: | |
| return img, gt | |
| elif pseudo is None and self.student_augs: | |
| return img, gt, aimg | |
| elif pseudo is not None and not self.student_augs: | |
| return img, gt, pseudo | |
| else: | |
| return img, gt, aimg, pseudo | |
| class ToTensor: | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[torch.Tensor, torch.Tensor]: | |
| img = F.to_tensor(np.array(img)) | |
| gt = torch.from_numpy(np.array(gt)).unsqueeze(0) | |
| if pseudo is not None: | |
| pseudo = torch.from_numpy(np.array(pseudo)).unsqueeze(0) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |
| class Resize: | |
| def __init__(self, | |
| resize: Tuple[int]): | |
| self.img_resize = tf.Resize(size=resize, | |
| interpolation=tf.InterpolationMode.BILINEAR) | |
| self.gt_resize = tf.Resize(size=resize, | |
| interpolation=tf.InterpolationMode.NEAREST) | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| img = self.img_resize(img) | |
| gt = self.gt_resize(gt) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, self.gt_resize(pseudo) | |
| class ImgResize: | |
| def __init__(self, | |
| resize: Tuple[int, int]): | |
| self.resize = resize | |
| self.num_pixels = self.resize[0]*self.resize[1] | |
| def __call__(self, | |
| img: torch.Tensor, | |
| gt: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| if torch.prod(torch.tensor(img.shape[-2:])) > self.num_pixels: | |
| img = torch.nn.functional.interpolate(img.unsqueeze(0), size=self.resize, mode='bilinear').squeeze(0) | |
| return img, gt | |
| class ImgResizePIL: | |
| def __init__(self, | |
| resize: Tuple[int]): | |
| self.resize = resize | |
| self.num_pixels = self.resize[0]*self.resize[1] | |
| def __call__(self, | |
| img: Image) -> Image: | |
| if img.height*img.width > self.num_pixels: | |
| img = img.resize((self.resize[1], self.resize[0]), tf.InterpolationMode.BILINEAR) | |
| return img | |
| class Normalize: | |
| def __init__(self, | |
| mean: List[float] = [0.485, 0.456, 0.406], | |
| std: List[float] = [0.229, 0.224, 0.225]): | |
| self.norm = tf.Normalize(mean=mean, | |
| std=std) | |
| def __call__(self, | |
| img: torch.Tensor, | |
| gt: torch.Tensor, | |
| pseudo = None) -> Tuple[torch.Tensor, torch.Tensor]: | |
| img = self.norm(img) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |
| class UnNormalize(object): | |
| def __init__(self, | |
| mean: List[float] = [0.485, 0.456, 0.406], | |
| std: List[float] = [0.229, 0.224, 0.225]): | |
| self.mean = mean | |
| self.std = std | |
| def __call__(self, image): | |
| image2 = torch.clone(image) | |
| for t, m, s in zip(image2, self.mean, self.std): | |
| t.mul_(s).add_(m) | |
| return image2 | |
| class RandomHFlip: | |
| def __init__(self, | |
| percentage: float = 0.5): | |
| self.percentage = percentage | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| if random.random() < self.percentage: | |
| img = F.hflip(img) | |
| gt = F.hflip(gt) | |
| if pseudo is not None: | |
| pseudo = F.hflip(pseudo) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |
| class RandomResizedCrop: | |
| def __init__(self, | |
| crop_size: List[int], | |
| crop_scale: List[float], | |
| crop_ratio: List[float]): | |
| print('RandomResizedCrop ratio modified!!!') | |
| self.crop_scale = tuple(crop_scale) | |
| self.crop_ratio = tuple(crop_ratio) | |
| self.crop = tf.RandomResizedCrop(size=tuple(crop_size), | |
| scale=self.crop_scale, | |
| ratio=self.crop_ratio,) | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| i, j, h, w = self.crop.get_params(img=img, | |
| scale=self.crop.scale, | |
| ratio=self.crop.ratio) | |
| img = F.resized_crop(img, i, j, h, w, self.crop.size, tf.InterpolationMode.BILINEAR) | |
| gt = F.resized_crop(gt, i, j, h, w, self.crop.size, tf.InterpolationMode.NEAREST) | |
| if pseudo is not None: | |
| pseudo = F.resized_crop(pseudo, i, j, h, w, self.crop.size, tf.InterpolationMode.NEAREST) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |
| class CenterCrop: | |
| def __init__(self, | |
| crop_size: int): | |
| self.crop = tf.CenterCrop(size=crop_size) | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| img = self.crop(img) | |
| gt = self.crop(gt) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, self.crop(pseudo) | |
| class PyramidCenterCrop: | |
| def __init__(self, | |
| crop_size: List[int], | |
| scales: List[float]): | |
| self.crop_size = crop_size | |
| self.scales = scales | |
| self.crop = tf.CenterCrop(size=crop_size) | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image) -> Tuple[Image.Image, Image.Image]: | |
| imgs = [] | |
| gts = [] | |
| for s in self.scales: | |
| new_size = (int(self.crop_size*1/s), int(self.crop_size*1/s*(img.shape[2]/img.shape[1]))) | |
| img = tf.Resize(size=new_size, interpolation=tf.InterpolationMode.BILINEAR)(img) | |
| gt = tf.Resize(size=new_size, interpolation=tf.InterpolationMode.NEAREST)(gt) | |
| imgs.append(self.crop(img)) | |
| gts.append(self.crop(gt)) | |
| return torch.stack(imgs), torch.stack(gts) | |
| class IdsToTrainIds: | |
| def __init__(self, | |
| source: str): | |
| self.source = source | |
| self.first_nonvoid = 7 | |
| def __call__(self, | |
| img: torch.Tensor, | |
| gt: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| if self.source == 'cityscapes': | |
| gt = gt.to(dtype=torch.int64) - self.first_nonvoid | |
| gt[gt>26] = 255 | |
| gt[gt<0] = 255 | |
| elif self.source == 'cocostuff': | |
| gt = gt.to(dtype=torch.int64) | |
| elif self.source == 'potsdam': | |
| gt = gt.to(dtype=torch.int64) | |
| return img, gt | |
| class ColorJitter: | |
| def __init__(self, percentage: float = 0.3, brightness: float = 0.1, | |
| contrast: float = 0.1, saturation: float = 0.1, hue: float = 0.1): | |
| self.percentage = percentage | |
| self.jitter = tf.ColorJitter(brightness=brightness, | |
| contrast=contrast, | |
| saturation=saturation, | |
| hue=hue) | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| if random.random() < self.percentage: | |
| img = self.jitter(img) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |
| class MaskGrayscale: | |
| def __init__(self, percentage: float = 0.1): | |
| self.percentage = percentage | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| if self.percentage > random.random(): | |
| img = tf.Grayscale(num_output_channels=3)(img) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |
| class RandGaussianBlur: | |
| def __init__(self, radius: List[float] = [.1, 2.]): | |
| self.radius = radius | |
| def __call__(self, | |
| img: Image.Image, | |
| gt: Image.Image, | |
| pseudo = None) -> Tuple[Image.Image, Image.Image]: | |
| radius = random.uniform(self.radius[0], self.radius[1]) | |
| img = tf.GaussianBlur(kernel_size=21, sigma=radius)(img) | |
| if pseudo is None: | |
| return img, gt | |
| else: | |
| return img, gt, pseudo | |