Spaces:
Paused
Paused
| """ basic augmentations | |
| """ | |
| import random | |
| import numpy as np | |
| import torch | |
| from torchvision import transforms | |
| import torch.nn.functional as F | |
| import torchvision.transforms.functional as TF | |
| import logging | |
| logger = logging.getLogger('root') | |
| def resize(sample, new_H, new_W): | |
| _, orig_H, orig_W = sample.img.shape | |
| sample.img = F.interpolate(sample.img.unsqueeze(0), size=(new_H, new_W), mode='bilinear', align_corners=False, antialias=True).squeeze(0) | |
| if sample.depth is not None: | |
| sample.depth = F.interpolate(sample.depth.unsqueeze(0), size=(new_H, new_W), mode='nearest').squeeze(0) | |
| if sample.depth_mask is not None: | |
| sample.depth_mask = F.interpolate(sample.depth_mask.unsqueeze(0).float(), size=(new_H, new_W), mode='nearest').squeeze(0) > 0.5 | |
| if sample.normal is not None: | |
| sample.normal = F.interpolate(sample.normal.unsqueeze(0), size=(new_H, new_W), mode='nearest').squeeze(0) | |
| if sample.normal_mask is not None: | |
| sample.normal_mask = F.interpolate(sample.normal_mask.unsqueeze(0).float(), size=(new_H, new_W), mode='nearest').squeeze(0) > 0.5 | |
| if sample.intrins is not None: | |
| # NOTE: top-left is (0,0) | |
| sample.intrins[0, 0] = sample.intrins[0, 0] * (new_W / orig_W) # fx | |
| sample.intrins[1, 1] = sample.intrins[1, 1] * (new_H / orig_H) # fy | |
| sample.intrins[0, 2] = (sample.intrins[0, 2] + 0.5) * (new_W / orig_W) - 0.5 # cx | |
| sample.intrins[1, 2] = (sample.intrins[1, 2] + 0.5) * (new_H / orig_H) - 0.5 # cy | |
| return sample | |
| def pad(sample, lrtb): | |
| l, r, t, b = lrtb | |
| sample.img = F.pad(sample.img, (l, r, t, b), mode="constant", value=0) | |
| if sample.depth is not None: | |
| sample.depth = F.pad(sample.depth, (l, r, t, b), mode="constant", value=0) | |
| if sample.depth_mask is not None: | |
| sample.depth_mask = F.pad(sample.depth_mask, (l, r, t, b), mode="constant", value=False) | |
| if sample.normal is not None: | |
| sample.normal = F.pad(sample.normal, (l, r, t, b), mode="constant", value=0) | |
| if sample.normal_mask is not None: | |
| sample.normal_mask = F.pad(sample.normal_mask, (l, r, t, b), mode="constant", value=False) | |
| if sample.intrins is not None: | |
| sample.intrins[0, 2] = sample.intrins[0, 2] + l | |
| sample.intrins[1, 2] = sample.intrins[1, 2] + t | |
| return sample | |
| def crop(sample, y, H, x, W): | |
| sample.img = sample.img[:, y:y+H, x:x+W] | |
| if sample.depth is not None: | |
| sample.depth = sample.depth[:, y:y+H, x:x+W] | |
| if sample.depth_mask is not None: | |
| sample.depth_mask = sample.depth_mask[:, y:y+H, x:x+W] | |
| if sample.normal is not None: | |
| sample.normal = sample.normal[:, y:y+H, x:x+W] | |
| if sample.normal_mask is not None: | |
| sample.normal_mask = sample.normal_mask[:, y:y+H, x:x+W] | |
| if sample.intrins is not None: | |
| sample.intrins[0, 2] = sample.intrins[0, 2] - x | |
| sample.intrins[1, 2] = sample.intrins[1, 2] - y | |
| return sample | |
| class ToTensor(): | |
| """ numpy arrays to torch tensors | |
| """ | |
| def __call__(self, sample): | |
| sample.img = torch.from_numpy(sample.img).permute(2, 0, 1) # (3, H, W) | |
| if sample.depth is not None: | |
| sample.depth = torch.from_numpy(sample.depth).permute(2, 0, 1) # (1, H, W) | |
| if sample.depth_mask is not None: | |
| sample.depth_mask = torch.from_numpy(sample.depth_mask).permute(2, 0, 1) # (1, H, W) | |
| if sample.normal is not None: | |
| sample.normal = torch.from_numpy(sample.normal).permute(2, 0, 1) # (3, H, W) | |
| if sample.normal_mask is not None: | |
| sample.normal_mask = torch.from_numpy(sample.normal_mask).permute(2, 0, 1) # (1, H, W) | |
| if sample.intrins is not None: | |
| sample.intrins = torch.from_numpy(sample.intrins) # (3, 3) | |
| return sample | |
| class RandomIntrins(): | |
| """ randomize intrinsics | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __call__(self, sample): | |
| assert 'crop_H' in sample.info.keys() | |
| assert 'crop_W' in sample.info.keys() | |
| crop_H = sample.info['crop_H'] | |
| crop_W = sample.info['crop_W'] | |
| # height-based resizing | |
| _, orig_H, orig_W = sample.img.shape | |
| new_H = random.randrange(min(orig_H, crop_H), max(orig_H, crop_H)+1) | |
| new_W = round((new_H / orig_H) * orig_W) | |
| sample = resize(sample, new_H=new_H, new_W=new_W) | |
| # pad if necessary | |
| orig_H, orig_W = sample.img.shape[1], sample.img.shape[2] | |
| l, r, t, b = 0, 0, 0, 0 | |
| if crop_H > orig_H: | |
| t = b = crop_H - orig_H | |
| if crop_W > orig_W: | |
| l = r = crop_W - orig_W | |
| sample = pad(sample, (l, r, t, b)) | |
| # crop | |
| assert sample.img.shape[1] >= crop_H | |
| assert sample.img.shape[2] >= crop_W | |
| x = random.randint(0, sample.img.shape[2] - crop_W) | |
| y = random.randint(0, sample.img.shape[1] - crop_H) | |
| sample = crop(sample, y=y, H=crop_H, x=x, W=crop_W) | |
| return sample | |
| class Resize(): | |
| """ resize to (H, W) | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __init__(self, H=480, W=640): | |
| self.H = H | |
| self.W = W | |
| def __call__(self, sample): | |
| return resize(sample, new_H=self.H, new_W=self.W) | |
| class RandomCrop(): | |
| """ random crop | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __init__(self, H=416, W=544): | |
| self.H = H | |
| self.W = W | |
| def __call__(self, sample): | |
| assert sample.img.shape[1] >= self.H | |
| assert sample.img.shape[2] >= self.W | |
| x = random.randint(0, sample.img.shape[2] - self.W) | |
| y = random.randint(0, sample.img.shape[1] - self.H) | |
| return crop(sample, y=y, H=self.H, x=x, W=self.W) | |
| class NyuCrop(): | |
| """ crop image border for NYUv2 images | |
| W = 43:608 / H = 45:472 | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __call__(self, sample): | |
| return crop(sample, y=45, H=472-45, x=43, W=608-43) | |
| class HorizontalFlip(): | |
| """ random horizontal flipping | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __init__(self, p=0.5): | |
| self.p = p | |
| def __call__(self, sample): | |
| if random.random() < self.p: | |
| sample.img = TF.hflip(sample.img) | |
| if sample.depth is not None: | |
| sample.depth = TF.hflip(sample.depth) | |
| if sample.depth_mask is not None: | |
| sample.depth_mask = TF.hflip(sample.depth_mask) | |
| if sample.normal is not None: | |
| sample.normal = TF.hflip(sample.normal) | |
| sample.normal[0, :, :] = -sample.normal[0, :, :] | |
| if sample.normal_mask is not None: | |
| sample.normal_mask = TF.hflip(sample.normal_mask) | |
| if sample.intrins is not None: | |
| # NOTE: top-left is (0,0) | |
| _, H, W = sample.img.shape | |
| sample.intrins[0, 2] = sample.intrins[0, 2] + 0.5 # top-left is (0.5, 0.5) | |
| sample.intrins[0, 2] = W - sample.intrins[0, 2] | |
| sample.intrins[0, 2] = sample.intrins[0, 2] - 0.5 # top-left is (0, 0) | |
| sample.flipped = True | |
| return sample | |
| class ColorAugmentation(): | |
| """ color augmentation | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __init__(self, gamma_range=(0.9, 1.1), | |
| brightness_range=(0.75, 1.25), | |
| color_range=(0.9, 1.1), | |
| p=0.5): | |
| self.gamma_range = gamma_range | |
| self.brightness_range = brightness_range | |
| self.color_range = color_range | |
| self.p = p | |
| def __call__(self, sample): | |
| if random.random() < self.p: | |
| # gamma augmentation | |
| gamma = random.uniform(*self.gamma_range) | |
| sample.img = sample.img ** gamma | |
| # brightness augmentation | |
| brightness = random.uniform(*self.brightness_range) | |
| sample.img = sample.img * brightness | |
| # color augmentation | |
| colors = np.random.uniform(*self.color_range, size=3).astype(np.float32) | |
| colors = torch.from_numpy(colors).view(3, 1, 1) | |
| sample.img = sample.img * colors | |
| # clip | |
| sample.img = torch.clip(sample.img, 0, 1) | |
| return sample | |
| class Normalize(): | |
| """ mean & std: for image normalization | |
| sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1] | |
| """ | |
| def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)): | |
| self.normalize = transforms.Normalize(mean=mean, std=std) | |
| def __call__(self, sample): | |
| sample.img = self.normalize(torch.clip(sample.img, min=0.0, max=1.0)) | |
| return sample | |
| class ToDict(): | |
| def __call__(self, sample): | |
| data_dict = {} | |
| for k, v in vars(sample).items(): | |
| if v is not None: | |
| data_dict[k] = v | |
| return data_dict | |