File size: 4,173 Bytes
079c32c |
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 |
from typing import Tuple, List
import random
import torch
import torch.nn as nn
class Intensity(nn.Module):
"""
Overview:
Intensity transformation for data augmentation. Scale the image intensity by a random factor.
"""
def __init__(self, scale: float) -> None:
"""
Arguments:
- scale (:obj:`float`): The scale factor for intensity transformation.
"""
super().__init__()
self.scale = scale
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Shapes:
- x (:obj:`torch.Tensor`): The input image tensor with shape (B, C, H, W).
- output (:obj:`torch.Tensor`): The output image tensor with shape (B, C, H, W).
"""
r = torch.randn((x.size(0), 1, 1, 1), device=x.device)
noise = 1.0 + (self.scale * r.clamp(-2.0, 2.0))
return x * noise
class RandomCrop(nn.Module):
"""
Overview:
Random crop the image to the given size.
"""
def __init__(self, image_shape: Tuple[int]) -> None:
"""
Arguments:
- image_shape (:obj:`Tuple[int]`): The target shape of the image to be cropped.
"""
super().__init__()
self.image_shape = image_shape
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Shapes:
- x (:obj:`torch.Tensor`): The input image tensor with shape (B, C, H, W), where H and W are \
the original image shape.
- output (:obj:`torch.Tensor`): The output image tensor with shape (B, C, H_, W_), where H_ and W_ are \
the target image shape indicated by `image_shape`.
"""
H, W = x.shape[2:]
H_, W_ = self.image_shape
dh, dw = H - H_, W - W_
h, w = random.randint(0, dh), random.randint(0, dw)
return x[..., h:h + H_, w:w + W_]
class ImageTransforms(object):
"""
Overview:
Image transformation for data augmentation. Including image normalization (divide 255), random crop and
intensity transformation.
"""
def __init__(self, augmentation: List[str], shift_delta: int = 4, image_shape: Tuple[int] = (96, 96)) -> None:
"""
Arguments:
- augmentation (:obj:`List[str]`): The list of augmentation types. Now support "shift" and "intensity".
- shift_delta (:obj:`int`): The delta value for random shift padding before crop. Use ReplicationPad2d \
to pad the image without the loss of information.
- image_shape (:obj:`Tuple[int]`): The target shape of the image to be cropped.
"""
self.augmentation = augmentation
self.image_transforms = []
for aug in self.augmentation:
if aug == "shift":
# TODO validate the effectiveness of ReflectionPad2d
transformation = nn.Sequential(nn.ReplicationPad2d(shift_delta), RandomCrop(image_shape))
elif aug == "intensity":
transformation = Intensity(scale=0.05)
else:
raise NotImplementedError("not support augmentation type: {}".format(aug))
self.image_transforms.append(transformation)
@torch.no_grad()
def transform(self, images: torch.Tensor) -> torch.Tensor:
"""
Shapes:
- x (:obj:`torch.Tensor`): The input image tensor with shape (B, C, H, W), where H and W are \
the original image shape.
- output (:obj:`torch.Tensor`): The output image tensor with shape (B, C, H_, W_), where H_ and W_ are \
the target image shape indicated by `image_shape`.
.. note::
Use torch.no_grad() to save cuda memory. Transformations are not trainable.
"""
images = images.float() / 255. if images.dtype == torch.uint8 else images
processed_images = images.reshape(-1, *images.shape[-3:])
for transform in self.image_transforms:
processed_images = transform(processed_images)
processed_images = processed_images.view(*images.shape[:-3], *processed_images.shape[1:])
return processed_images
|