from matplotlib import pyplot as plt
import cv2
import numpy as np
import torch
import kornia
import torchvision
#############################
# We use OpenCV to load an image to memory represented in a numpy.array
img_bgr: np.array = cv2.imread('./data/drslump.jpg', cv2.IMREAD_COLOR)
#############################
# Convert the numpy array to torch
x_bgr: torch.Tensor = kornia.image_to_tensor(img_bgr)
x_rgb: torch.Tensor = kornia.bgr_to_rgb(x_bgr)
#############################
# Create batch and normalize
x_rgb = x_rgb.expand(2, -1, -1, -1) # 4xCxHxW
x_rgb = x_rgb.float() / 255.
def imshow(input: torch.Tensor):
out: torch.Tensor = torchvision.utils.make_grid(input, nrow=2, padding=1)
out_np: np.array = kornia.tensor_to_image(out)
plt.imshow(out_np)
plt.axis('off')
#############################
return torch.flip(input, [-2, -1])
def imshow(input: torch.Tensor):
out: torch.Tensor = torchvision.utils.make_grid(input, nrow=2, padding=5)
out_np: np.ndarray = kornia.tensor_to_image(out)
plt.imshow(out_np)
plt.axis('off')
#############################
# Create a batch of images
xb_bgr = torch.cat([x_bgr, hflip(x_bgr), vflip(x_bgr), rot180(x_bgr)])
imshow(xb_bgr)
#############################
# Convert BGR to RGB
xb_rgb = kornia.bgr_to_rgb(xb_bgr)
imshow(xb_rgb)
#############################
# Convert RGB to grayscale
# NOTE: image comes in torch.uint8, and kornia assumes floating point type
xb_gray = kornia.rgb_to_grayscale(xb_rgb.float() / 255.)
imshow(xb_gray)
#############################
# Convert RGB to HSV
xb_hsv = kornia.rgb_to_hsv(xb_rgb.float() / 255.)
imshow(xb_hsv[:, 2:3])
def test_back_and_forth(self, device):
data_bgr = torch.rand(1, 3, 3, 2).to(device)
data_rgb = kornia.bgr_to_rgb(data_bgr)
data_bgr_new = kornia.rgb_to_bgr(data_rgb)
assert_allclose(data_bgr, data_bgr_new)
def op_script(data: torch.Tensor) -> torch.Tensor:
return kornia.bgr_to_rgb(data)
import numpy as np
import torch
import kornia
#############################
# We use OpenCV to load an image to memory represented in a numpy.ndarray
img_bgr: np.ndarray = cv2.imread('./data/arturito.jpeg') # HxWxC
#############################
# The image is convert to a 4D torch tensor
x_bgr: torch.tensor = kornia.image_to_tensor(img_bgr) # 1xCxHxW
#############################
# Once with a torch tensor we can use any Kornia operator
x_rgb: torch.tensor = kornia.bgr_to_rgb(x_bgr) # 1xCxHxW
#############################
# Convert back to numpy to visualize
img_rgb: np.ndarray = kornia.tensor_to_image(x_rgb.byte()) # HxWxC
#############################
# We use Matplotlib to visualize de results
fig, axs = plt.subplots(1, 2, figsize=(32, 16))
axs = axs.ravel()
axs[0].axis('off')
axs[0].imshow(img_bgr)
axs[1].axis('off')
axs[1].imshow(img_rgb)
def __getitem__(self, item):
img = cv2.imread(self.arg.dataset_route[self.dataset] +
self.list[item][-1])
img = image_to_tensor(img)
img = bgr_to_rgb(img)
return img
def __getitem__(self, item):
dataset_route, dataset, split, type, annotation, crop_size, RGB, sigma, trans_ratio, rotate_limit,\
scale_ratio_up, scale_ratio_down, scale_horizontal, scale_vertical =\
self.arg.dataset_route, self.dataset, self.split, self.type,\
self.list[item], self.arg.crop_size, self.arg.RGB, self.arg.sigma,\
self.arg.trans_ratio, self.arg.rotate_limit, self.arg.scale_ratio_up, self.arg.scale_ratio_down,\
self.arg.scale_horizontal, self.arg.scale_vertical
pic_orig = cv2.imread(dataset_route[dataset] + annotation[-1])
coord_x = list(map(float, annotation[:2 * kp_num[dataset]:2]))
coord_y = list(map(float, annotation[1:2 * kp_num[dataset]:2]))
bbox = np.array(list(map(int, annotation[-7:-3])))
translation, trans_dir, rotation, scaling, scaling_horizontal, scaling_vertical, flip, gaussian_blur = get_random_transform_param(
type,
bbox,
trans_ratio,
rotate_limit,
scale_ratio_up,
scale_ratio_down,
scale_horizontal,
scale_vertical,
flip=False,
gaussian=False)
horizontal_add = (bbox[2] - bbox[0]) * (1 - scaling)
vertical_add = (bbox[3] - bbox[1]) * (1 - scaling)
bbox = np.float32([
bbox[0] - horizontal_add, bbox[1] - vertical_add,
bbox[2] + horizontal_add, bbox[3] + vertical_add
])
horizontal_add = (bbox[2] - bbox[0]) * scaling_horizontal
vertical_add = (bbox[3] - bbox[1]) * scaling_vertical
bbox = np.float32([
bbox[0] - horizontal_add, bbox[1] - vertical_add,
bbox[2] + horizontal_add, bbox[3] + vertical_add
])
position_before = np.float32(
[[
int(bbox[0]) + pow(-1, trans_dir + 1) * translation,
int(bbox[1]) + pow(-1, trans_dir // 2 + 1) * translation
],
[
int(bbox[0]) + pow(-1, trans_dir + 1) * translation,
int(bbox[3]) + pow(-1, trans_dir // 2 + 1) * translation
],
[
int(bbox[2]) + pow(-1, trans_dir + 1) * translation,
int(bbox[3]) + pow(-1, trans_dir // 2 + 1) * translation
]])
position_after = np.float32([[0, 0], [0, crop_size - 1],
[crop_size - 1, crop_size - 1]])
crop_matrix = cv2.getAffineTransform(position_before, position_after)
# crop_matrix = np.vstack([crop_matrix, [0, 0, 1]])
pic_affine_orig = cv2.warpAffine(pic_orig,
crop_matrix, (crop_size, crop_size),
borderMode=cv2.BORDER_REPLICATE)
# width_height = (bbox[2] - bbox[0], bbox[3] - bbox[1])
width_height = (crop_size, crop_size)
affine_matrix = get_affine_matrix(width_height, rotation, scaling)
# affine_matrix = np.vstack([affine_matrix, [0, 0, 1]])
# affine_matrix = np.matmul(crop_matrix, affine_matrix)
# TODO one transform
pic_affine_orig = cv2.warpAffine(pic_affine_orig,
affine_matrix, (crop_size, crop_size),
borderMode=cv2.BORDER_REPLICATE)
pic_affine_orig = further_transform(
pic_affine_orig, bbox, flip,
gaussian_blur) if type in ['train'] else pic_affine_orig
# show_img(pic_affine_orig, wait=0, keep=False)
pic_affine_orig = bgr_to_rgb(image_to_tensor(pic_affine_orig))
pic_affine_orig_norm = normalize(pic_affine_orig,
torch.from_numpy(self.mean_color),
torch.from_numpy(self.std_color))
if not RGB:
pic_affine = convert_img_to_gray(pic_affine_orig)
pic_affine = normalize(pic_affine, self.mean_gray, self.std_gray)
else:
pic_affine = pic_affine_orig_norm
coord_x_cropped, coord_y_cropped = get_cropped_coords(dataset,
crop_matrix,
coord_x,
coord_y,
crop_size,
flip=flip)
gt_coords_xy = get_gt_coords(dataset, affine_matrix, coord_x_cropped,
coord_y_cropped)
return pic_affine, pic_affine_orig_norm, gt_coords_xy