pythonpytorchtorchvisiondata-augmentation

Rotation with Pytorch Torchvision but from two specific angles

I want to randomly rotate my images at 45 degrees or -45 degrees, but only these two, no in between. 50% chances for 45 and 50% for -45.

I know that I can use RandomRotation of torchvision, but if I do like in this example:

import torchvision.transforms as T

transforms = T.Compose([
  T.RandomRotation((-45, 45))
])

The rotation will be random between -45 and 45 degrees, so that's not what I want. How to be sure the rotation will be specifically these two angles?

Thanks in advance for your answers.

Solution

  • Like Julien said, just make a 50/50 random choice, then do the fixed angle rotation. Below is an example

    import torchvision.transforms as T

transforms = T.Compose([
    T.RandomChoice([
        T.RandomRotation((45, 45)),
        T.RandomRotation((-45, -45)),
    ], p=[0.5, 0.5])
])