import random
import torch
import torch.nn.functional as F
def apply_cr_aug(x, flip=True, translation=True):
if flip:
x = random_flip(x, 0.5)
if translation:
x = random_translation(x, 1 / 8)
if flip or translation:
x = x.contiguous()
return x
def random_flip(x, p):
x_out = x.clone()
n, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
flip_prob = torch.FloatTensor(n, 1).uniform_(0.0, 1.0).to(x.device)
flip_mask = flip_prob < p
flip_mask = flip_mask.type(torch.bool).view(n, 1, 1, 1).repeat(1, c, h, w)
x_out[flip_mask] = torch.flip(x[flip_mask].view(-1, c, h, w), [3]).view(-1)
return x_out
def random_translation(x, ratio):
max_t_x, max_t_y = int(x.shape[2] * ratio), int(x.shape[3] * ratio)
t_x = torch.randint(-max_t_x, max_t_x + 1, size=[x.shape[0], 1, 1], device=x.device)
t_y = torch.randint(-max_t_y, max_t_y + 1, size=[x.shape[0], 1, 1], device=x.device)
grid_batch, grid_x, grid_y = torch.meshgrid(
torch.arange(x.shape[0], dtype=torch.long, device=x.device),
torch.arange(x.shape[2], dtype=torch.long, device=x.device),
torch.arange(x.shape[3], dtype=torch.long, device=x.device),
)
grid_x = (grid_x + t_x) + max_t_x
grid_y = (grid_y + t_y) + max_t_y
x_pad = F.pad(input=x, pad=[max_t_y, max_t_y, max_t_x, max_t_x], mode='reflect')
x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
return x
