Vincentqyw
fix: roma
c74a070
raw
history blame
12.8 kB
# From https://github.com/TRI-ML/KP2D.
# Copyright 2020 Toyota Research Institute. All rights reserved.
import random
from math import pi
import cv2
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from utils import image_grid
def filter_dict(dict, keywords):
"""
Returns only the keywords that are part of a dictionary
Parameters
----------
dictionary : dict
Dictionary for filtering
keywords : list of str
Keywords that will be filtered
Returns
-------
keywords : list of str
List containing the keywords that are keys in dictionary
"""
return [key for key in keywords if key in dict]
def resize_sample(sample, image_shape, image_interpolation=Image.ANTIALIAS):
"""
Resizes a sample, which contains an input image.
Parameters
----------
sample : dict
Dictionary with sample values (output from a dataset's __getitem__ method)
shape : tuple (H,W)
Output shape
image_interpolation : int
Interpolation mode
Returns
-------
sample : dict
Resized sample
"""
# image
image_transform = transforms.Resize(image_shape, interpolation=image_interpolation)
sample["image"] = image_transform(sample["image"])
return sample
def spatial_augment_sample(sample):
"""Apply spatial augmentation to an image (flipping and random affine transformation)."""
augment_image = transforms.Compose(
[
transforms.RandomVerticalFlip(p=0.5),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomAffine(15, translate=(0.1, 0.1), scale=(0.9, 1.1)),
]
)
sample["image"] = augment_image(sample["image"])
return sample
def unnormalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
"""Counterpart method of torchvision.transforms.Normalize."""
for t, m, s in zip(tensor, mean, std):
t.div_(1 / s).sub_(-m)
return tensor
def sample_homography(
shape,
perspective=True,
scaling=True,
rotation=True,
translation=True,
n_scales=100,
n_angles=100,
scaling_amplitude=0.1,
perspective_amplitude=0.4,
patch_ratio=0.8,
max_angle=pi / 4,
):
"""Sample a random homography that includes perspective, scale, translation and rotation operations."""
width = float(shape[1])
hw_ratio = float(shape[0]) / float(shape[1])
pts1 = np.stack([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]], axis=0)
pts2 = pts1.copy() * patch_ratio
pts2[:, 1] *= hw_ratio
if perspective:
perspective_amplitude_x = np.random.normal(0.0, perspective_amplitude / 2, (2))
perspective_amplitude_y = np.random.normal(
0.0, hw_ratio * perspective_amplitude / 2, (2)
)
perspective_amplitude_x = np.clip(
perspective_amplitude_x,
-perspective_amplitude / 2,
perspective_amplitude / 2,
)
perspective_amplitude_y = np.clip(
perspective_amplitude_y,
hw_ratio * -perspective_amplitude / 2,
hw_ratio * perspective_amplitude / 2,
)
pts2[0, 0] -= perspective_amplitude_x[1]
pts2[0, 1] -= perspective_amplitude_y[1]
pts2[1, 0] -= perspective_amplitude_x[0]
pts2[1, 1] += perspective_amplitude_y[1]
pts2[2, 0] += perspective_amplitude_x[1]
pts2[2, 1] -= perspective_amplitude_y[0]
pts2[3, 0] += perspective_amplitude_x[0]
pts2[3, 1] += perspective_amplitude_y[0]
if scaling:
random_scales = np.random.normal(1, scaling_amplitude / 2, (n_scales))
random_scales = np.clip(
random_scales, 1 - scaling_amplitude / 2, 1 + scaling_amplitude / 2
)
scales = np.concatenate([[1.0], random_scales], 0)
center = np.mean(pts2, axis=0, keepdims=True)
scaled = (
np.expand_dims(pts2 - center, axis=0)
* np.expand_dims(np.expand_dims(scales, 1), 1)
+ center
)
valid = np.arange(n_scales) # all scales are valid except scale=1
idx = valid[np.random.randint(valid.shape[0])]
pts2 = scaled[idx]
if translation:
t_min, t_max = np.min(pts2 - [-1.0, -hw_ratio], axis=0), np.min(
[1.0, hw_ratio] - pts2, axis=0
)
pts2 += np.expand_dims(
np.stack(
[
np.random.uniform(-t_min[0], t_max[0]),
np.random.uniform(-t_min[1], t_max[1]),
]
),
axis=0,
)
if rotation:
angles = np.linspace(-max_angle, max_angle, n_angles)
angles = np.concatenate([[0.0], angles], axis=0)
center = np.mean(pts2, axis=0, keepdims=True)
rot_mat = np.reshape(
np.stack(
[np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)],
axis=1,
),
[-1, 2, 2],
)
rotated = (
np.matmul(
np.tile(np.expand_dims(pts2 - center, axis=0), [n_angles + 1, 1, 1]),
rot_mat,
)
+ center
)
valid = np.where(
np.all(
(rotated >= [-1.0, -hw_ratio]) & (rotated < [1.0, hw_ratio]),
axis=(1, 2),
)
)[0]
idx = valid[np.random.randint(valid.shape[0])]
pts2 = rotated[idx]
pts2[:, 1] /= hw_ratio
def ax(p, q):
return [p[0], p[1], 1, 0, 0, 0, -p[0] * q[0], -p[1] * q[0]]
def ay(p, q):
return [0, 0, 0, p[0], p[1], 1, -p[0] * q[1], -p[1] * q[1]]
a_mat = np.stack([f(pts1[i], pts2[i]) for i in range(4) for f in (ax, ay)], axis=0)
p_mat = np.transpose(
np.stack([[pts2[i][j] for i in range(4) for j in range(2)]], axis=0)
)
homography = np.matmul(np.linalg.pinv(a_mat), p_mat).squeeze()
homography = np.concatenate([homography, [1.0]]).reshape(3, 3)
return homography
def warp_homography(sources, homography):
"""Warp features given a homography
Parameters
----------
sources: torch.tensor (1,H,W,2)
Keypoint vector.
homography: torch.Tensor (3,3)
Homography.
Returns
-------
warped_sources: torch.tensor (1,H,W,2)
Warped feature vector.
"""
_, H, W, _ = sources.shape
warped_sources = sources.clone().squeeze()
warped_sources = warped_sources.view(-1, 2)
warped_sources = torch.addmm(
homography[:, 2], warped_sources, homography[:, :2].t()
)
warped_sources.mul_(1 / warped_sources[:, 2].unsqueeze(1))
warped_sources = warped_sources[:, :2].contiguous().view(1, H, W, 2)
return warped_sources
def add_noise(img, mode="gaussian", percent=0.02):
"""Add image noise
Parameters
----------
image : np.array
Input image
mode: str
Type of noise, from ['gaussian','salt','pepper','s&p']
percent: float
Percentage image points to add noise to.
Returns
-------
image : np.array
Image plus noise.
"""
original_dtype = img.dtype
if mode == "gaussian":
mean = 0
var = 0.1
sigma = var * 0.5
if img.ndim == 2:
h, w = img.shape
gauss = np.random.normal(mean, sigma, (h, w))
else:
h, w, c = img.shape
gauss = np.random.normal(mean, sigma, (h, w, c))
if img.dtype not in [np.float32, np.float64]:
gauss = gauss * np.iinfo(img.dtype).max
img = np.clip(img.astype(np.float) + gauss, 0, np.iinfo(img.dtype).max)
else:
img = np.clip(img.astype(np.float) + gauss, 0, 1)
elif mode == "salt":
print(img.dtype)
s_vs_p = 1
num_salt = np.ceil(percent * img.size * s_vs_p)
coords = tuple([np.random.randint(0, i - 1, int(num_salt)) for i in img.shape])
if img.dtype in [np.float32, np.float64]:
img[coords] = 1
else:
img[coords] = np.iinfo(img.dtype).max
print(img.dtype)
elif mode == "pepper":
s_vs_p = 0
num_pepper = np.ceil(percent * img.size * (1.0 - s_vs_p))
coords = tuple(
[np.random.randint(0, i - 1, int(num_pepper)) for i in img.shape]
)
img[coords] = 0
elif mode == "s&p":
s_vs_p = 0.5
# Salt mode
num_salt = np.ceil(percent * img.size * s_vs_p)
coords = tuple([np.random.randint(0, i - 1, int(num_salt)) for i in img.shape])
if img.dtype in [np.float32, np.float64]:
img[coords] = 1
else:
img[coords] = np.iinfo(img.dtype).max
# Pepper mode
num_pepper = np.ceil(percent * img.size * (1.0 - s_vs_p))
coords = tuple(
[np.random.randint(0, i - 1, int(num_pepper)) for i in img.shape]
)
img[coords] = 0
else:
raise ValueError("not support mode for {}".format(mode))
noisy = img.astype(original_dtype)
return noisy
def non_spatial_augmentation(
img_warp_ori, jitter_paramters, color_order=[0, 1, 2], to_gray=False
):
"""Apply non-spatial augmentation to an image (jittering, color swap, convert to gray scale, Gaussian blur)."""
brightness, contrast, saturation, hue = jitter_paramters
color_augmentation = transforms.ColorJitter(brightness, contrast, saturation, hue)
"""
augment_image = color_augmentation.get_params(brightness=[max(0, 1 - brightness), 1 + brightness],
contrast=[max(0, 1 - contrast), 1 + contrast],
saturation=[max(0, 1 - saturation), 1 + saturation],
hue=[-hue, hue])
"""
B = img_warp_ori.shape[0]
img_warp = []
kernel_sizes = [0, 1, 3, 5]
for b in range(B):
img_warp_sub = img_warp_ori[b].cpu()
img_warp_sub = torchvision.transforms.functional.to_pil_image(img_warp_sub)
img_warp_sub_np = np.array(img_warp_sub)
img_warp_sub_np = img_warp_sub_np[:, :, color_order]
if np.random.rand() > 0.5:
img_warp_sub_np = add_noise(img_warp_sub_np)
rand_index = np.random.randint(4)
kernel_size = kernel_sizes[rand_index]
if kernel_size > 0:
img_warp_sub_np = cv2.GaussianBlur(
img_warp_sub_np, (kernel_size, kernel_size), sigmaX=0
)
if to_gray:
img_warp_sub_np = cv2.cvtColor(img_warp_sub_np, cv2.COLOR_RGB2GRAY)
img_warp_sub_np = cv2.cvtColor(img_warp_sub_np, cv2.COLOR_GRAY2RGB)
img_warp_sub = Image.fromarray(img_warp_sub_np)
img_warp_sub = color_augmentation(img_warp_sub)
img_warp_sub = torchvision.transforms.functional.to_tensor(img_warp_sub).to(
img_warp_ori.device
)
img_warp.append(img_warp_sub)
img_warp = torch.stack(img_warp, dim=0)
return img_warp
def ha_augment_sample(
data,
jitter_paramters=[0.5, 0.5, 0.2, 0.05],
patch_ratio=0.7,
scaling_amplitude=0.2,
max_angle=pi / 4,
):
"""Apply Homography Adaptation image augmentation."""
input_img = data["image"].unsqueeze(0)
_, _, H, W = input_img.shape
device = input_img.device
homography = (
torch.from_numpy(
sample_homography(
[H, W],
patch_ratio=patch_ratio,
scaling_amplitude=scaling_amplitude,
max_angle=max_angle,
)
)
.float()
.to(device)
)
homography_inv = torch.inverse(homography)
source = (
image_grid(
1, H, W, dtype=input_img.dtype, device=device, ones=False, normalized=True
)
.clone()
.permute(0, 2, 3, 1)
)
target_warped = warp_homography(source, homography)
img_warp = torch.nn.functional.grid_sample(input_img, target_warped)
color_order = [0, 1, 2]
if np.random.rand() > 0.5:
random.shuffle(color_order)
to_gray = False
if np.random.rand() > 0.5:
to_gray = True
input_img = non_spatial_augmentation(
input_img,
jitter_paramters=jitter_paramters,
color_order=color_order,
to_gray=to_gray,
)
img_warp = non_spatial_augmentation(
img_warp,
jitter_paramters=jitter_paramters,
color_order=color_order,
to_gray=to_gray,
)
data["image"] = input_img.squeeze()
data["image_aug"] = img_warp.squeeze()
data["homography"] = homography
data["homography_inv"] = homography_inv
return data