File size: 691 Bytes
7667a87 | 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 | from typing import Sequence
import random
from torch import nn
import torchvision.transforms.functional as TF
class TransformsFixRotation(nn.Module):
r"""
Rotate by one of the given angles.
Example: `rotation_transform = MyRotationTransform(angles=[-30, -15, 0, 15, 30])`
"""
def __init__(self, angles):
super(TransformsFixRotation, self).__init__()
if not isinstance(angles, Sequence):
angles = [angles, ]
self.angles = angles
def forward(self, x):
angle = random.choice(self.angles)
return TF.rotate(x, angle)
def __repr__(self) -> str:
return f"{self.__class__.__name__}(angles={self.angles})"
|