Spaces:
Sleeping
Sleeping
File size: 2,996 Bytes
6df18f5 |
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 |
import torch
from torch import nn
import kornia.augmentation as K
class ImageAugmentations(nn.Module):
def __init__(self, output_size, augmentations_number, p=0.7, resize=True):
super().__init__()
self.output_size = output_size
self.augmentations_number = augmentations_number
self.augmentations = [
K.RandomAffine(degrees=15, translate=0.1, p=p, padding_mode="border"),
# K.RandomPerspective(0.7, p=p),
]
self.resize = (
nn.AdaptiveAvgPool2d((self.output_size, self.output_size))
if resize
else (lambda x: x)
)
def forward(self, image, mask, with_orig=True):
"""Extends the image and mask with identical augmentations
If the input consists of image I, and mask M, the extended augmented output will be:
[I_aug1, I_aug2, I_aug3, ...], [M_aug1, M_aug2, M_aug3, ...]
If with_orig=True, the extended augmented output will be:
[I, I_aug1, I_aug2, ...], [M, M_aug1, M_aug2, ...]
Args:
image: tensor of shape [1, C, H, W]
mask: tensor of shape [1, 1, H, W]
with_orig: if True, first returned image and mask will be un-augmented inputs
Returns:
tuple of (extended images of shape [augmentations_number, C, H, W],
extended masks of shape [augmentations_number, 1, H, W])
"""
# Duplicate the inputs, in contrast to regular augmentations that do not change the number of samples
resized_images = self.resize(image)
resized_images = resized_images.repeat(self.augmentations_number, 1, 1, 1)
resized_masks = self.resize(mask)
resized_masks = resized_masks.repeat(self.augmentations_number, 1, 1, 1)
batch_size = image.shape[0]
if with_orig:
# At least one non-augmented image
non_aug_inputs = resized_images[:batch_size]
aug_inputs = resized_images[batch_size:]
non_aug_masks = resized_masks[:batch_size]
aug_masks = resized_masks[batch_size:]
for trans in self.augmentations:
trans_params = trans.forward_parameters(aug_inputs.shape)
aug_inputs = trans(aug_inputs, trans_params)
aug_masks = trans(aug_masks, trans_params)
updated_input_batch = torch.cat([non_aug_inputs, aug_inputs], dim=0)
updated_mask_batch = torch.cat([non_aug_masks, aug_masks], dim=0)
else:
aug_inputs = resized_images
aug_masks = resized_masks
for trans in self.augmentations:
trans_params = trans.forward_parameters(aug_inputs.shape)
aug_inputs = trans(aug_inputs, trans_params)
aug_masks = trans(aug_masks, trans_params)
updated_input_batch = aug_inputs
updated_mask_batch = aug_masks
return updated_input_batch, updated_mask_batch
|