def test_rotate(self):
x = np.zeros((100, 100, 3), dtype=np.uint8)
x[40, 40] = [255, 255, 255]
with self.assertRaises(TypeError):
F.rotate(x, 10)
img = F.to_pil_image(x)
result = F.rotate(img, 45)
assert result.size == (100, 100)
r, c, ch = np.where(result)
assert all(x in r for x in [49, 50])
assert all(x in c for x in [36])
assert all(x in ch for x in [0, 1, 2])
result = F.rotate(img, 45, expand=True)
assert result.size == (142, 142)
r, c, ch = np.where(result)
assert all(x in r for x in [70, 71])
assert all(x in c for x in [57])
assert all(x in ch for x in [0, 1, 2])
result = F.rotate(img, 45, center=(40, 40))
assert result.size == (100, 100)
r, c, ch = np.where(result)
assert all(x in r for x in [40])
assert all(x in c for x in [40])
assert all(x in ch for x in [0, 1, 2])
result_a = F.rotate(img, 90)
result_b = F.rotate(img, -270)
assert np.all(np.array(result_a) == np.array(result_b))
def tf_rotation_stack(x, num_rotations=8):
xs = []
for i in range(num_rotations):
angle = 360 * i / num_rotations
xrot = TF.rotate(x, angle)
xrot = TF.to_tensor(xrot)
xs.append(xrot)
xs = torch.stack(xs)
return xs
def __call__(self, *args):
angle = random.uniform(self.min, self.max)
rst = []
for pic in args:
rst.append(
functional.rotate(pic, angle, self.resample, self.expand,
self.center, self.fill))
return rst
def __call__(self, image: torch.Tensor,
context: ExpressionContext) -> torch.Tensor:
degree = context(self.degree)
center = context(self.center)
center_pix = [
int(center[0] * image.shape[-1]),
int(center[1] * image.shape[-2]),
]
return VF.rotate(image, degree, center=center_pix)
def __call__(self, img):
angle = self.get_params(self.degrees)
for i, im in enumerate(img):
img[i] = F.rotate(im, angle, self.resample, self.expand,
self.center)
return img
def __call__(self, img, *args, **kwargs):
if random.uniform(0, 1) > self.p:
return img
rnd_angle = random.randint(self.angle[0], self.angle[1])
return F.rotate(img,
rnd_angle,
resample=False,
expand=False,
center=None)
def __call__(self, sample):
"""
img (PIL Image): Image to be rotated.
Returns:
PIL Image: Rotated image.
"""
angle = self.get_params(self.degrees)
sample['image'] = F.rotate(sample['image'],
angle=angle,
resample=Image.BICUBIC,
center=self.center)
sample['label'] = F.rotate(sample['label'],
angle=angle,
resample=Image.NEAREST,
center=self.center)
return sample
def test_rotate(self):
# Tests on square image
scripted_rotate = torch.jit.script(F.rotate)
data = [
self._create_data(26, 26, device=self.device),
self._create_data(32, 26, device=self.device)
]
for tensor, pil_img in data:
img_size = pil_img.size
centers = [
None, (int(img_size[0] * 0.3), int(img_size[0] * 0.4)),
[int(img_size[0] * 0.5),
int(img_size[0] * 0.6)]
]
for r in [
0,
]:
for a in range(-180, 180, 17):
for e in [True, False]:
for c in centers:
out_pil_img = F.rotate(pil_img,
angle=a,
resample=r,
expand=e,
center=c)
out_pil_tensor = torch.from_numpy(
np.array(out_pil_img).transpose((2, 0, 1)))
for fn in [F.rotate, scripted_rotate]:
out_tensor = fn(tensor,
angle=a,
resample=r,
expand=e,
center=c).cpu()
self.assertEqual(out_tensor.shape,
out_pil_tensor.shape,
msg="{}: {} vs {}".format(
(img_size, r, a, e, c),
out_tensor.shape,
out_pil_tensor.shape))
num_diff_pixels = (out_tensor != out_pil_tensor
).sum().item() / 3.0
ratio_diff_pixels = num_diff_pixels / out_tensor.shape[
-1] / out_tensor.shape[-2]
# Tolerance : less than 2% of different pixels
self.assertLess(
ratio_diff_pixels,
0.02,
msg="{}: {}\n{} vs \n{}".format(
(img_size, r, a, e, c),
ratio_diff_pixels,
out_tensor[0, :7, :7],
out_pil_tensor[0, :7, :7]))
def transform(image, mask, image_size=224):
# Resize
resized_num = int(random() * image_size)
resize = transforms.Resize(size=(image_size + resized_num, image_size + resized_num))
image = resize(image)
mask = resize(mask)
# num_pad = int(random.random() * image_size)
# image = TF.pad(image, num_pad, padding_mode='edge')
# mask = TF.pad(mask, num_pad)
resize = transforms.Resize(size=(image_size, image_size))
image = resize(image)
mask = resize(mask)
# Make gray scale image
gray_image = TF.to_grayscale(image)
# Data Augmentation
if random() > 0.5:
image = TF.vflip(image)
gray_image = TF.vflip(gray_image)
mask = TF.vflip(mask)
if random() > 0.5:
image = TF.hflip(image)
gray_image = TF.hflip(gray_image)
mask = TF.hflip(mask)
if random() > 0.5:
angle = random() * 12 - 6
image = TF.rotate(image, angle)
gray_image = TF.rotate(gray_image, angle)
mask = TF.rotate(mask, angle)
# Transform to tensor
image = TF.to_tensor(image)
mask = TF.to_tensor(mask)[0]
gray_image = TF.to_tensor(gray_image)
# Normalize Data
image = TF.normalize(image, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
return image, gray_image, mask
def rotate_sample(sample, degrees=20):
"""
Rotates input images as data augmentation.
Assumes the intrinsics to be scaled w.r.t the image. i.e. this step should be followed by (image &intrinsics) resizing
Parameters
----------
sample : dict
Input sample
degrees : float
sample a random rotation angle between [-degrees,, degrees]
Returns
-------
sample : dict
Rotated sample
Author: Zeeshan Khan Suri
"""
if degrees == 0:
return sample
# Do same rotation transform for image and context imgs
rand_degree = (torch.rand(1) - 0.5) * 2 * degrees
# Get rotation center as principal point from intrinsic matrix
center = sample["intrinsics"][:2, 2]
# Jitter single items
for key in filter_dict(sample, ['rgb', 'rgb_original', 'depth']):
sample[key] = TF.rotate(
sample[key],
rand_degree.item(),
resample=Image.NEAREST if key == 'depth' else Image.BILINEAR,
center=center.tolist())
# Jitter lists
for key in filter_dict(
sample, ['rgb_context', 'rgb_context_original', 'depth_context']):
sample[key] = [
TF.rotate(k,
rand_degree.item(),
resample=Image.NEAREST
if key == 'depth_context' else Image.BILINEAR,
center=center.tolist()) for k in sample[key]
]
# Return rotated (?) sample
return sample
def transform(img_input, img_target, opticalflow_1, opticalflow_2, patchsize):
opticalflow_1 = opticalflow_1.permute(2,0,1)
opticalflow_2 = opticalflow_2.permute(2,0,1)
# Random rotation
angle = transforms.RandomRotation.get_params([-10, 10])
img_input = TF.rotate(img_input, angle) #pytorch:1.7.0, torchvision:0.8.1
img_target = TF.rotate(img_target, angle)
opticalflow_1 = TF.rotate(opticalflow_1, angle)
opticalflow_2 = TF.rotate(opticalflow_2, angle)
# First crop to (520,1080) to eliminate black margin zone
img_input = transforms.CenterCrop((520,1080))(img_input)
img_target = transforms.CenterCrop((520,1080))(img_target)
opticalflow_1 = transforms.CenterCrop((520,1080))(opticalflow_1)
opticalflow_2 = transforms.CenterCrop((520,1080))(opticalflow_2)
# Random crop to (256,256)
i, j, h, w = transforms.RandomCrop.get_params(
img_input, output_size=(patchsize, patchsize))
img_input = TF.crop(img_input, i, j, h, w)
img_target = TF.crop(img_target, i, j, h, w)
opticalflow_1 = TF.crop(opticalflow_1, i, j, h, w)
opticalflow_2 = TF.crop(opticalflow_2, i, j, h, w)
# Random horizontal flipping
if random.random() > 0.5:
img_input = TF.hflip(img_input)
img_target = TF.hflip(img_target)
opticalflow_1 = TF.hflip(opticalflow_1)
opticalflow_2 = TF.hflip(opticalflow_2)
# Random vertical flipping
if random.random() > 0.5:
img_input = TF.vflip(img_input)
img_target = TF.vflip(img_target)
opticalflow_1 = TF.vflip(opticalflow_1)
opticalflow_2 = TF.vflip(opticalflow_2)
opticalflow_1 = transforms.Resize((patchsize//2, patchsize//2))(opticalflow_1)
opticalflow_2 = transforms.Resize((patchsize//2, patchsize//2))(opticalflow_2)
return img_input, img_target, opticalflow_1, opticalflow_2
def pairedTransformations(self, img, gtImg, config, mode):
# Resize images
img = TF.to_pil_image(img)
gtImg = TF.to_pil_image(gtImg)
img = TF.resize(img,
size=(config["imgSize"], config["imgSize"]),
interpolation=2)
gtImg = TF.resize(gtImg,
size=(config["imgSize"], config["imgSize"]),
interpolation=0)
if (self.mode == 'train'):
# Rotate images
if (random.choice([True, False])
and 'Rotate' in config['augmentations']):
angle = random.randint(-10, -10)
img = TF.rotate(img, angle)
gtImg = TF.rotate(gtImg, angle)
# Randomly crop images
if (random.choice([True, False])
and 'Crop' in config['augmentations']):
i, j, h, w = transforms.RandomResizedCrop.get_params(
gtImg, scale=(0.8, 1), ratio=(0.75, 1))
img = TF.resized_crop(img,
i,
j,
h,
w,
size=(config["imgSize"],
config["imgSize"]),
interpolation=2)
gtImg = TF.resized_crop(gtImg,
i,
j,
h,
w,
size=(config["imgSize"],
config["imgSize"]),
interpolation=0)
# Standardize
img = TF.to_tensor(img)
img = TF.normalize(img, mean=[img.mean()], std=[img.std()])
gtImg = torch.from_numpy(np.expand_dims(np.array(gtImg), 0))
return img, gtImg
def train_data_augmentation(self, p, x, y):
dim_x = x.shape[-1]
dim_y = y.shape[-1]
mask_size = x.shape[0] + 12
new_x = torch.empty((dim_x, mask_size, mask_size))
new_y = torch.empty((dim_y, mask_size, mask_size))
hflip, vflip = False, False
angle = 0
top, left = 0, 0
if np.random.rand() > p:
hflip = True
if np.random.rand() > p:
vflip = True
if np.random.rand() > 0:
angle = np.random.randint(-30, 30)
if np.random.rand() > p:
top, left = np.random.randint(0, 8), np.random.randint(0, 8)
for i in range(dim_x):
tmp_x = x[:, :, i:i + 1]
tmp_img = tf.to_pil_image((tmp_x))
tmp_img = tf.pad(tmp_img, 6)
if hflip:
tmp_img = tf.hflip(tmp_img)
if vflip:
tmp_img = tf.vflip(tmp_img)
tmp_img = tf.pad(tmp_img, 4)
tmp_img = tf.crop(tmp_img, top, left, mask_size, mask_size)
tmp_img = tmp_img.convert('P')
tmp_img = tf.rotate(tmp_img, angle)
new_x[i:i + 1, :, :] = tf.to_tensor(tmp_img)
for i in range(dim_y):
tmp_y = y[:, :, i:i + 1]
tmp_img = tf.to_pil_image((tmp_y))
tmp_img = tf.pad(tmp_img, 6)
if hflip:
tmp_img = tf.hflip(tmp_img)
if vflip:
tmp_img = tf.vflip(tmp_img)
tmp_img = tf.pad(tmp_img, 4)
tmp_img = tf.crop(tmp_img, top, left, mask_size, mask_size)
tmp_img = tmp_img.convert('P')
tmp_img = tf.rotate(tmp_img, angle)
new_y[i:i + 1, :, :] = tf.to_tensor(tmp_img)
return new_x, new_y
def __call__(self, image, mask):
image = self.to_tensor(image, normalize=True,
repeat_channels=3).float()
mask = self.to_tensor(mask, normalize=False).type(torch.LongTensor).to(
self._device)
if self.new_size:
image = TF.resize(image, size=list(self.new_size))
mask = TF.resize(mask,
size=list(self.new_size),
interpolation=TF.InterpolationMode.NEAREST)
if self.mode == "train":
angle = random.choice(self.angles)
image = TF.rotate(image, angle)
mask = TF.rotate(mask, angle)
# vertical flip
if random.random() < self.p_flip:
image = TF.vflip(image)
mask = TF.vflip(mask)
# horizontal flip
if random.random() < self.p_flip:
image = TF.hflip(image)
mask = TF.hflip(mask)
if random.random() < self.p_crop:
# Resize
resize = T.Resize(size=(480, 672),
interpolation=TF.InterpolationMode.NEAREST)
# Random crop
i, j, h, w = T.RandomCrop.get_params(image,
output_size=(320, 448))
image = TF.crop(image, i, j, h, w)
mask = TF.crop(mask, i, j, h, w)
image = resize(image)
mask = resize(mask)
# Contrast
return image, mask
def transform_augmentation(self, image, flip, rotation, affine_angle, affine_shear):
# Horizontal flip
if flip:
image = functional.hflip(image)
# Rotate image.
image = functional.rotate(image, rotation)
# Affine transformation
# image = functional.affine(image, affine_angle, (0,0), affine_shear) # Affine not available in this pytorch version
return image
def cart2polar(images_cart, channels, timesteps):
m, C, H, W = images_cart.shape
mid_pt = int(H // 2)
R = torch.linspace(0, mid_pt, channels).long()
thetas = torch.linspace(0, 360, timesteps).float()
images_polar = []
for theta in thetas:
image_rotated = TF.rotate(images_cart, theta.item())
images_polar.append(image_rotated[:, :, mid_pt, R])
return torch.cat(images_polar, axis=1).transpose(1, 2)
def __call__(self, x: np.ndarray) -> np.ndarray:
"""
Applies the rotation.
:param x: image to be rotated
:return: rotated image
"""
degs = (self.iteration * self.increase_per_iteration +
self.degrees) % 360
self.iteration += 1
return F.rotate(x, degs)
def __call__(self, sample):
angle = self.get_params(self.degrees)
rdict = {}
input_data = sample['input']
input_data = [
F.rotate(input, angle, self.resample, self.expand, self.center)
for input in input_data
]
rdict['input'] = input_data
gt_data = sample['gt']
gt_data = [
F.rotate(gt, angle, self.resample, self.expand, self.center)
for gt in gt_data
]
rdict['gt'] = gt_data
return rdict
def create_rotated_images_and_labels(images):
images = images.cpu()
images = [tf.to_pil_image(x) for x in images]
number_of_images = len(images)
images = [tf.rotate(x, 0) for x in images] \
+ [tf.rotate(x, 90) for x in images] \
+ [tf.rotate(x, 180) for x in images] \
+ [tf.rotate(x, 270) for x in images]
rotation_labels = np.repeat(0, number_of_images).tolist() \
+ np.repeat(1, number_of_images).tolist() \
+ np.repeat(2, number_of_images).tolist() \
+ np.repeat(3, number_of_images).tolist()
images = [tf.to_tensor(x) for x in images]
images = torch.stack(images).to(DEVICE)
labels = torch.LongTensor(rotation_labels).to(DEVICE)
return images, labels
def __call__(self, img):
"""
img (PIL Image): Image to be rotated.
Returns:
PIL Image: Rotated image.
"""
angle = self.get_params(self.degrees)
return F.rotate(img, angle, self.resample, self.expand, self.center)
def __call__(self, *imgs):
angle = self.get_angle()
if angle == int(angle) and angle % 90 == 0:
if angle == 0:
return imgs
else:
#print(imgs)
return map(lambda x: F.rotate(x, angle, False, False, None), imgs)
else:
return map(lambda x: self._pad_rotate(x, angle), imgs)
def my_transform(self, image, mask):
# Resize
if self.train and self.scale:
min_size = int(self.size_img[0] * self.scale_factor[0])
max_size = int(self.size_img[0] * self.scale_factor[1])
if min_size < self.size_crop[0]:
size = random.randint(self.size_crop[0], max_size)
else:
size = random.randint(min_size, max_size)
resize = T.Resize((size, size))
else:
resize = T.Resize(self.size_img)
image = resize(image)
mask = resize(mask)
if self.train:
# Random crop
i, j, h, w = T.RandomCrop.get_params(image,
output_size=self.size_crop)
image = TF.crop(image, i, j, h, w)
mask = TF.crop(mask, i, j, h, w)
# Random horizontal flipping
if random.random() > self.p:
image = TF.hflip(image)
mask = TF.hflip(mask)
if self.rotate:
if random.random() > self.p_rotate:
if random.random() > 0.5:
angle = np.random.randint(0, 30)
image = TF.rotate(image, angle=angle)
mask = TF.rotate(mask, angle=angle)
else:
angle = np.random.randint(330, 360)
image = TF.rotate(image, angle=angle)
mask = TF.rotate(mask, angle=angle)
# Transform to tensor
image = TF.to_tensor(image)
image = TF.normalize(image, self.mean, self.std)
mask = to_tensor_target(mask)
return image, mask
def apply_augmentations(self, img, target):
# Horizontal flip
if random.random() > 0.5:
img = TF.hflip(img)
target = TF.hflip(target)
# Random Rotation (clockwise and counter clockwise)
if random.random() > 0.5:
degrees = 10
if random.random() > 0.5:
degrees *= -1
img = TF.rotate(img, degrees)
target = TF.rotate(target, degrees)
# Brighten or darken image (only applied to input image)
if random.random() > 0.5:
brightness = 1.2
if random.random() > 0.5:
brightness -= 0.4
img = TF.adjust_brightness(img, brightness)
return img, target
def __getitem__(self, index):
# image_name = self.image_arr[index] # Get image name from csv
image_label = self.label_arr[index] # Get image label from csv
img_as_img = self._loaded_images[index]
angle = random.choice(
[0, 90]
) # Don't need -90, because that is achieved in combination with the flips.
img_as_img = TF.rotate(img_as_img, angle)
img_as_tensor = self.to_tensor(img_as_img) # Already scaled to [0,1]
return img_as_tensor, image_label
def __getitem__(self, index):
image = cv2.imread(
os.path.join(self.img_dir, self.list[index] + '.png'))
label = np.load(os.path.join(self.label_dir,
self.list[index] + '.npy'))
height, width = label.shape
if self.train and self.augment:
# random rotations
if np.random.randint(2) == 0:
ang = np.random.choice([90, -90])
image = np.dstack([
F.rotate(self._np2pil(image[:, :, i]), ang)
for i in range(3)
])
label = np.asarray(F.rotate(self._np2pil(label), ang))
# random h-flips
if np.random.randint(2) == 0:
image = np.dstack(
[F.hflip(self._np2pil(image[:, :, i])) for i in range(3)])
label = np.asarray(F.hflip(self._np2pil(label)))
# random v-flips
if np.random.randint(2) == 0:
image = np.dstack(
[F.vflip(self._np2pil(image[:, :, i])) for i in range(3)])
label = np.asarray(F.vflip(self._np2pil(label)))
# random crops
if np.random.randint(2) == 0:
i, j, h, w = transforms.RandomCrop.get_params(
self._np2pil(label), output_size=(height // 2, width // 2))
image = np.dstack([
F.resized_crop(self._np2pil(image[:, :, ii]), i, j, h, w,
(height, width)) for ii in range(3)
])
label = np.asarray(
F.resized_crop(self._np2pil(label), i, j, h, w,
(height, width)))
return self._to_tensor(image), self._to_tensor(label)
def _aug_img(self, image):
if random.random() > 0.5:
image = TF.rotate(image, random.choice([90, 180, 270]))
'''
if random.random() > 0.5:
image = TF.hflip(image)
if random.random() > 0.5:
image = TF.vflip(image)
'''
return image
def train_rotation_self_supervised(self, batch, step):
batch_image_4x = []
label_4x = []
for image in batch['sketch_img']:
batch_image_4x.append(image)
label_4x.append(0)
image = TF.normalize(
image,
mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
std=[1 / 0.229, 1 / 0.224, 1 / 0.225])
pil_image = TF.to_pil_image(image)
tensor_rotate = TF.normalize(TF.to_tensor(TF.rotate(pil_image,
90)),
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
batch_image_4x.append(tensor_rotate)
label_4x.append(1)
tensor_rotate = TF.normalize(TF.to_tensor(TF.rotate(
pil_image, 180)),
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
batch_image_4x.append(tensor_rotate)
label_4x.append(2)
tensor_rotate = TF.normalize(TF.to_tensor(TF.rotate(
pil_image, 270)),
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
batch_image_4x.append(tensor_rotate)
label_4x.append(3)
random_indices = torch.randperm(len(label_4x))
batch_image_4x = torch.stack(batch_image_4x)[random_indices]
label_4x = torch.tensor(label_4x)[random_indices]
self.train()
self.step = step
self.optimizer.zero_grad()
output = self.Network(batch_image_4x.to(device))
loss = self.loss(output, label_4x.to(device))
loss.backward()
self.optimizer.step()
return loss.item()
def __call__(self, img):
"""
Args:
img (PIL Image): Image to be flipped.
Returns:
PIL Image: Randomly flipped image.
"""
if np.random.rand() < self.p:
return F.rotate(img, 90)
return img
def random_rotate(self, image, label, angle=None):
'''
:param image: PIL RGB uint8
:param label: PIL, uint8
:param angle: None, list-float, tuple-float
:return: PIL
'''
if angle is None:
angle = transforms.RandomRotation.get_params([-180, 180])
elif isinstance(angle, list) or isinstance(angle, tuple):
angle = random.choice(angle)
image = tf.rotate(image, angle, fill=self.image_fill)
label = tf.rotate(label, angle, fill=self.label_fill)
image = tf.resize(image, self.input_hw, interpolation=Image.BILINEAR)
label = tf.resize(label, self.input_hw, interpolation=Image.NEAREST)
return image, label
def __call__(self, x, y=None):
if len(x.shape) != 3:
raise ValueError(
"Input of RandomAxialRotation3D should be a 3 dimensionnal tensor."
)
angle = self.get_params(self.degrees)
x_rotated = np.zeros(x.shape, dtype=x.dtype)
y_rotated = np.zeros(y.shape, dtype=y.dtype) if y is not None else None
for i in range(x.shape[0]):
x_rotated[i, :, :] = F.rotate(
Image.fromarray(x[i, :, :], mode='F'), angle)
if y is not None:
y_rotated[x, :, :] = F.rotate(
Image.fromarray(y[i, :, :], mode='F'), angle)
if y is not None:
return x, y
return x
def __getitem__(self, index):
sample_path = self.samples[index]
img = Image.open(sample_path).convert('RGB')
hr = self.cropper(img)
lr = hr.copy()
lr = self.resizer(lr)
if self.rotate and np.random.rand() < 0.5:
rv = np.random.randint(1, 4)
lr = TF.rotate(lr, 90 * rv)
hr = TF.rotate(hr, 90 * rv)
if self.hflip and np.random.rand() < 0.5:
lr = TF.hflip(lr)
hr = TF.hflip(hr)
if self.vflip and np.random.rand() < 0.5:
lr = TF.vflip(lr)
hr = TF.vflip(hr)
return TF.to_tensor(lr), TF.to_tensor(hr)
def torchvision(self, img):
return torchvision.rotate(img, angle=-45, resample=Image.BILINEAR)
