def apply_transformation(x, trans):
dx, dy, angle = trans[0], trans[1], trans[2]
height, width = x.shape[2], x.shape[3]
# Pad the image to prevent two-step rotation / translation from truncating
# corners
max_dist_from_center = np.sqrt(height**2 + width**2) / 2
min_edge_from_center = float(np.min([height, width])) / 2
padding = np.ceil(max_dist_from_center - min_edge_from_center).astype(
np.int32)
x = nn.ConstantPad2d(padding, 0)(x)
# Apply rotation
angle = ch.from_numpy(np.ones(x.shape[0]) * angle)
angle = angle.to(x.get_device())
x = rotate(x, angle)
# Apply translation
dx_in_px = -dx * height
dy_in_px = -dy * width
translation = ch.from_numpy(
np.tile(np.array([dx_in_px, dy_in_px], dtype=np.float32),
(x.shape[0], 1)))
translation = translation.to(x.get_device())
x = translate(x, translation)
x = translate(x, translation)
# Pad if needed
if x.shape[2] < height or x.shape[3] < width:
pad = nn.ConstantPad2d(
(0, max(0, height - x.shape[2]), 0, max(0, width - x.shape[3])), 0)
x = pad(x)
return center_crop(x, (height, width))
def inner(image_t):
dx = np.random.choice(d)
dy = np.random.choice(d)
return translate(image_t, torch.tensor([[dx, dy]]).float().to(device))
def translate_y(x: Tensor, v: float) -> Tensor:
B, C, H, W = x.shape
return T.translate(
x, torch.tensor([[0, v * H]], device=x.device, dtype=x.dtype))