mirror of
https://github.com/aljazceru/InvSR.git
synced 2025-12-17 06:14:22 +01:00
1150 lines
39 KiB
Python
1150 lines
39 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding:utf-8 -*-
|
|
# Power by Zongsheng Yue 2021-11-24 16:54:19
|
|
|
|
import sys
|
|
import cv2
|
|
import math
|
|
import torch
|
|
import random
|
|
import numpy as np
|
|
from scipy import fft
|
|
from pathlib import Path
|
|
from einops import rearrange
|
|
from skimage import img_as_ubyte, img_as_float32
|
|
|
|
# --------------------------Metrics----------------------------
|
|
def ssim(img1, img2):
|
|
C1 = (0.01 * 255)**2
|
|
C2 = (0.03 * 255)**2
|
|
|
|
img1 = img1.astype(np.float64)
|
|
img2 = img2.astype(np.float64)
|
|
kernel = cv2.getGaussianKernel(11, 1.5)
|
|
window = np.outer(kernel, kernel.transpose())
|
|
|
|
mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid
|
|
mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
|
|
mu1_sq = mu1**2
|
|
mu2_sq = mu2**2
|
|
mu1_mu2 = mu1 * mu2
|
|
sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
|
|
sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
|
|
sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
|
|
|
|
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
|
|
(sigma1_sq + sigma2_sq + C2))
|
|
return ssim_map.mean()
|
|
|
|
def calculate_ssim(im1, im2, border=0, ycbcr=False):
|
|
'''
|
|
SSIM the same outputs as MATLAB's
|
|
im1, im2: h x w x , [0, 255], uint8
|
|
'''
|
|
if not im1.shape == im2.shape:
|
|
raise ValueError('Input images must have the same dimensions.')
|
|
|
|
if ycbcr:
|
|
im1 = rgb2ycbcr(im1, True)
|
|
im2 = rgb2ycbcr(im2, True)
|
|
|
|
h, w = im1.shape[:2]
|
|
im1 = im1[border:h-border, border:w-border]
|
|
im2 = im2[border:h-border, border:w-border]
|
|
|
|
if im1.ndim == 2:
|
|
return ssim(im1, im2)
|
|
elif im1.ndim == 3:
|
|
if im1.shape[2] == 3:
|
|
ssims = []
|
|
for i in range(3):
|
|
ssims.append(ssim(im1[:,:,i], im2[:,:,i]))
|
|
return np.array(ssims).mean()
|
|
elif im1.shape[2] == 1:
|
|
return ssim(np.squeeze(im1), np.squeeze(im2))
|
|
else:
|
|
raise ValueError('Wrong input image dimensions.')
|
|
|
|
def calculate_psnr(im1, im2, border=0, ycbcr=False):
|
|
'''
|
|
PSNR metric.
|
|
im1, im2: h x w x , [0, 255], uint8
|
|
'''
|
|
if not im1.shape == im2.shape:
|
|
raise ValueError('Input images must have the same dimensions.')
|
|
|
|
if ycbcr:
|
|
im1 = rgb2ycbcr(im1, True)
|
|
im2 = rgb2ycbcr(im2, True)
|
|
|
|
h, w = im1.shape[:2]
|
|
im1 = im1[border:h-border, border:w-border]
|
|
im2 = im2[border:h-border, border:w-border]
|
|
|
|
im1 = im1.astype(np.float64)
|
|
im2 = im2.astype(np.float64)
|
|
mse = np.mean((im1 - im2)**2)
|
|
if mse == 0:
|
|
return float('inf')
|
|
return 20 * math.log10(255.0 / math.sqrt(mse))
|
|
|
|
def normalize_np(im, mean=0.5, std=0.5, reverse=False):
|
|
'''
|
|
Input:
|
|
im: h x w x c, numpy array
|
|
Normalize: (im - mean) / std
|
|
Reverse: im * std + mean
|
|
|
|
'''
|
|
if not isinstance(mean, (list, tuple)):
|
|
mean = [mean, ] * im.shape[2]
|
|
mean = np.array(mean).reshape([1, 1, im.shape[2]])
|
|
|
|
if not isinstance(std, (list, tuple)):
|
|
std = [std, ] * im.shape[2]
|
|
std = np.array(std).reshape([1, 1, im.shape[2]])
|
|
|
|
if not reverse:
|
|
out = (im.astype(np.float32) - mean) / std
|
|
else:
|
|
out = im.astype(np.float32) * std + mean
|
|
return out
|
|
|
|
def normalize_th(im, mean=0.5, std=0.5, reverse=False):
|
|
'''
|
|
Input:
|
|
im: b x c x h x w, torch tensor
|
|
Normalize: (im - mean) / std
|
|
Reverse: im * std + mean
|
|
|
|
'''
|
|
if not isinstance(mean, (list, tuple)):
|
|
mean = [mean, ] * im.shape[1]
|
|
mean = torch.tensor(mean, device=im.device).view([1, im.shape[1], 1, 1])
|
|
|
|
if not isinstance(std, (list, tuple)):
|
|
std = [std, ] * im.shape[1]
|
|
std = torch.tensor(std, device=im.device).view([1, im.shape[1], 1, 1])
|
|
|
|
if not reverse:
|
|
out = (im - mean) / std
|
|
else:
|
|
out = im * std + mean
|
|
return out
|
|
|
|
# ------------------------Image format--------------------------
|
|
def rgb2ycbcr(im, only_y=True):
|
|
'''
|
|
same as matlab rgb2ycbcr
|
|
Input:
|
|
im: uint8 [0,255] or float [0,1]
|
|
only_y: only return Y channel
|
|
'''
|
|
# transform to float64 data type, range [0, 255]
|
|
if im.dtype == np.uint8:
|
|
im_temp = im.astype(np.float64)
|
|
else:
|
|
im_temp = (im * 255).astype(np.float64)
|
|
|
|
# convert
|
|
if only_y:
|
|
rlt = np.dot(im_temp, np.array([65.481, 128.553, 24.966])/ 255.0) + 16.0
|
|
else:
|
|
rlt = np.matmul(im_temp, np.array([[65.481, -37.797, 112.0 ],
|
|
[128.553, -74.203, -93.786],
|
|
[24.966, 112.0, -18.214]])/255.0) + [16, 128, 128]
|
|
if im.dtype == np.uint8:
|
|
rlt = rlt.round()
|
|
else:
|
|
rlt /= 255.
|
|
return rlt.astype(im.dtype)
|
|
|
|
def rgb2ycbcrTorch(im, only_y=True):
|
|
'''
|
|
same as matlab rgb2ycbcr
|
|
Input:
|
|
im: float [0,1], N x 3 x H x W
|
|
only_y: only return Y channel
|
|
'''
|
|
# transform to range [0,255.0]
|
|
im_temp = im.permute([0,2,3,1]) * 255.0 # N x H x W x C --> N x H x W x C
|
|
# convert
|
|
if only_y:
|
|
rlt = torch.matmul(im_temp, torch.tensor([65.481, 128.553, 24.966],
|
|
device=im.device, dtype=im.dtype).view([3,1])/ 255.0) + 16.0
|
|
else:
|
|
scale = torch.tensor(
|
|
[[65.481, -37.797, 112.0 ],
|
|
[128.553, -74.203, -93.786],
|
|
[24.966, 112.0, -18.214]],
|
|
device=im.device, dtype=im.dtype
|
|
) / 255.0
|
|
bias = torch.tensor([16, 128, 128], device=im.device, dtype=im.dtype).view([-1, 1, 1, 3])
|
|
rlt = torch.matmul(im_temp, scale) + bias
|
|
|
|
rlt /= 255.0
|
|
rlt.clamp_(0.0, 1.0)
|
|
return rlt.permute([0, 3, 1, 2])
|
|
|
|
def ycbcr2rgbTorch(im):
|
|
'''
|
|
same as matlab ycbcr2rgb
|
|
Input:
|
|
im: float [0,1], N x 3 x H x W
|
|
only_y: only return Y channel
|
|
'''
|
|
# transform to range [0,255.0]
|
|
im_temp = im.permute([0,2,3,1]) * 255.0 # N x H x W x C --> N x H x W x C
|
|
# convert
|
|
scale = torch.tensor(
|
|
[[0.00456621, 0.00456621, 0.00456621],
|
|
[0, -0.00153632, 0.00791071],
|
|
[0.00625893, -0.00318811, 0]],
|
|
device=im.device, dtype=im.dtype
|
|
) * 255.0
|
|
bias = torch.tensor(
|
|
[-222.921, 135.576, -276.836], device=im.device, dtype=im.dtype
|
|
).view([-1, 1, 1, 3])
|
|
rlt = torch.matmul(im_temp, scale) + bias
|
|
rlt /= 255.0
|
|
rlt.clamp_(0.0, 1.0)
|
|
return rlt.permute([0, 3, 1, 2])
|
|
|
|
def bgr2rgb(im): return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
|
|
|
def rgb2bgr(im): return cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
|
|
|
|
def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
|
|
"""Convert torch Tensors into image numpy arrays.
|
|
|
|
After clamping to [min, max], values will be normalized to [0, 1].
|
|
|
|
Args:
|
|
tensor (Tensor or list[Tensor]): Accept shapes:
|
|
1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
|
|
2) 3D Tensor of shape (3/1 x H x W);
|
|
3) 2D Tensor of shape (H x W).
|
|
Tensor channel should be in RGB order.
|
|
rgb2bgr (bool): Whether to change rgb to bgr.
|
|
out_type (numpy type): output types. If ``np.uint8``, transform outputs
|
|
to uint8 type with range [0, 255]; otherwise, float type with
|
|
range [0, 1]. Default: ``np.uint8``.
|
|
min_max (tuple[int]): min and max values for clamp.
|
|
|
|
Returns:
|
|
(Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
|
|
shape (H x W). The channel order is BGR.
|
|
"""
|
|
if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
|
|
raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
|
|
|
|
flag_tensor = torch.is_tensor(tensor)
|
|
if flag_tensor:
|
|
tensor = [tensor]
|
|
result = []
|
|
for _tensor in tensor:
|
|
_tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
|
|
_tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
|
|
|
|
n_dim = _tensor.dim()
|
|
if n_dim == 4:
|
|
img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
|
|
img_np = img_np.transpose(1, 2, 0)
|
|
if rgb2bgr:
|
|
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
|
elif n_dim == 3:
|
|
img_np = _tensor.numpy()
|
|
img_np = img_np.transpose(1, 2, 0)
|
|
if img_np.shape[2] == 1: # gray image
|
|
img_np = np.squeeze(img_np, axis=2)
|
|
else:
|
|
if rgb2bgr:
|
|
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
|
elif n_dim == 2:
|
|
img_np = _tensor.numpy()
|
|
else:
|
|
raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
|
|
if out_type == np.uint8:
|
|
# Unlike MATLAB, numpy.unit8() WILL NOT round by default.
|
|
img_np = (img_np * 255.0).round()
|
|
img_np = img_np.astype(out_type)
|
|
result.append(img_np)
|
|
if len(result) == 1 and flag_tensor:
|
|
result = result[0]
|
|
return result
|
|
|
|
def img2tensor(imgs, bgr2rgb=False, out_type=torch.float32):
|
|
"""Convert image numpy arrays into torch tensor.
|
|
Args:
|
|
imgs (Array or list[array]): Accept shapes:
|
|
3) list of numpy arrays
|
|
1) 3D numpy array of shape (H x W x 3/1);
|
|
2) 2D Tensor of shape (H x W).
|
|
Tensor channel should be in RGB order.
|
|
|
|
Returns:
|
|
(array or list): 4D ndarray of shape (1 x C x H x W)
|
|
"""
|
|
|
|
def _img2tensor(img):
|
|
if img.ndim == 2:
|
|
tensor = torch.from_numpy(img[None, None,]).type(out_type)
|
|
elif img.ndim == 3:
|
|
if bgr2rgb:
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
tensor = torch.from_numpy(rearrange(img, 'h w c -> c h w')).type(out_type).unsqueeze(0)
|
|
else:
|
|
raise TypeError(f'2D or 3D numpy array expected, got{img.ndim}D array')
|
|
return tensor
|
|
|
|
if not (isinstance(imgs, np.ndarray) or (isinstance(imgs, list) and all(isinstance(t, np.ndarray) for t in imgs))):
|
|
raise TypeError(f'Numpy array or list of numpy array expected, got {type(imgs)}')
|
|
|
|
flag_numpy = isinstance(imgs, np.ndarray)
|
|
if flag_numpy:
|
|
imgs = [imgs,]
|
|
result = []
|
|
for _img in imgs:
|
|
result.append(_img2tensor(_img))
|
|
|
|
if len(result) == 1 and flag_numpy:
|
|
result = result[0]
|
|
return result
|
|
|
|
# ------------------------Image resize-----------------------------
|
|
def imresize_np(img, scale, antialiasing=True):
|
|
# Now the scale should be the same for H and W
|
|
# input: img: Numpy, HWC or HW [0,1]
|
|
# output: HWC or HW [0,1] w/o round
|
|
img = torch.from_numpy(img)
|
|
need_squeeze = True if img.dim() == 2 else False
|
|
if need_squeeze:
|
|
img.unsqueeze_(2)
|
|
|
|
in_H, in_W, in_C = img.size()
|
|
out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
|
|
kernel_width = 4
|
|
kernel = 'cubic'
|
|
|
|
# Return the desired dimension order for performing the resize. The
|
|
# strategy is to perform the resize first along the dimension with the
|
|
# smallest scale factor.
|
|
# Now we do not support this.
|
|
|
|
# get weights and indices
|
|
weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
|
|
in_H, out_H, scale, kernel, kernel_width, antialiasing)
|
|
weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
|
|
in_W, out_W, scale, kernel, kernel_width, antialiasing)
|
|
# process H dimension
|
|
# symmetric copying
|
|
img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
|
|
img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
|
|
|
|
sym_patch = img[:sym_len_Hs, :, :]
|
|
inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
|
|
sym_patch_inv = sym_patch.index_select(0, inv_idx)
|
|
img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
|
|
|
|
sym_patch = img[-sym_len_He:, :, :]
|
|
inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
|
|
sym_patch_inv = sym_patch.index_select(0, inv_idx)
|
|
img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
|
|
|
|
out_1 = torch.FloatTensor(out_H, in_W, in_C)
|
|
kernel_width = weights_H.size(1)
|
|
for i in range(out_H):
|
|
idx = int(indices_H[i][0])
|
|
for j in range(out_C):
|
|
out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i])
|
|
|
|
# process W dimension
|
|
# symmetric copying
|
|
out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
|
|
out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
|
|
|
|
sym_patch = out_1[:, :sym_len_Ws, :]
|
|
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
|
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
|
out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
|
|
|
|
sym_patch = out_1[:, -sym_len_We:, :]
|
|
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
|
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
|
out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
|
|
|
|
out_2 = torch.FloatTensor(out_H, out_W, in_C)
|
|
kernel_width = weights_W.size(1)
|
|
for i in range(out_W):
|
|
idx = int(indices_W[i][0])
|
|
for j in range(out_C):
|
|
out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i])
|
|
if need_squeeze:
|
|
out_2.squeeze_()
|
|
|
|
return out_2.numpy()
|
|
|
|
def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
|
|
if (scale < 1) and (antialiasing):
|
|
# Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
|
|
kernel_width = kernel_width / scale
|
|
|
|
# Output-space coordinates
|
|
x = torch.linspace(1, out_length, out_length)
|
|
|
|
# Input-space coordinates. Calculate the inverse mapping such that 0.5
|
|
# in output space maps to 0.5 in input space, and 0.5+scale in output
|
|
# space maps to 1.5 in input space.
|
|
u = x / scale + 0.5 * (1 - 1 / scale)
|
|
|
|
# What is the left-most pixel that can be involved in the computation?
|
|
left = torch.floor(u - kernel_width / 2)
|
|
|
|
# What is the maximum number of pixels that can be involved in the
|
|
# computation? Note: it's OK to use an extra pixel here; if the
|
|
# corresponding weights are all zero, it will be eliminated at the end
|
|
# of this function.
|
|
P = math.ceil(kernel_width) + 2
|
|
|
|
# The indices of the input pixels involved in computing the k-th output
|
|
# pixel are in row k of the indices matrix.
|
|
indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
|
|
1, P).expand(out_length, P)
|
|
|
|
# The weights used to compute the k-th output pixel are in row k of the
|
|
# weights matrix.
|
|
distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
|
|
# apply cubic kernel
|
|
if (scale < 1) and (antialiasing):
|
|
weights = scale * cubic(distance_to_center * scale)
|
|
else:
|
|
weights = cubic(distance_to_center)
|
|
# Normalize the weights matrix so that each row sums to 1.
|
|
weights_sum = torch.sum(weights, 1).view(out_length, 1)
|
|
weights = weights / weights_sum.expand(out_length, P)
|
|
|
|
# If a column in weights is all zero, get rid of it. only consider the first and last column.
|
|
weights_zero_tmp = torch.sum((weights == 0), 0)
|
|
if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
|
|
indices = indices.narrow(1, 1, P - 2)
|
|
weights = weights.narrow(1, 1, P - 2)
|
|
if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
|
|
indices = indices.narrow(1, 0, P - 2)
|
|
weights = weights.narrow(1, 0, P - 2)
|
|
weights = weights.contiguous()
|
|
indices = indices.contiguous()
|
|
sym_len_s = -indices.min() + 1
|
|
sym_len_e = indices.max() - in_length
|
|
indices = indices + sym_len_s - 1
|
|
return weights, indices, int(sym_len_s), int(sym_len_e)
|
|
|
|
# matlab 'imresize' function, now only support 'bicubic'
|
|
def cubic(x):
|
|
absx = torch.abs(x)
|
|
absx2 = absx**2
|
|
absx3 = absx**3
|
|
return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
|
|
(-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
|
|
|
|
# ------------------------Image I/O-----------------------------
|
|
def imread(path, chn='rgb', dtype='float32', force_gray2rgb=True, force_rgba2rgb=False):
|
|
'''
|
|
Read image.
|
|
chn: 'rgb', 'bgr' or 'gray'
|
|
out:
|
|
im: h x w x c, numpy tensor
|
|
'''
|
|
try:
|
|
im = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) # BGR, uint8
|
|
except:
|
|
print(str(path))
|
|
|
|
if im is None:
|
|
print(str(path))
|
|
|
|
if chn.lower() == 'gray':
|
|
assert im.ndim == 2, f"{str(path)} can't be successfuly read!"
|
|
else:
|
|
if im.ndim == 2:
|
|
if force_gray2rgb:
|
|
im = np.stack([im, im, im], axis=2)
|
|
else:
|
|
raise ValueError(f"{str(path)} has {im.ndim} channels!")
|
|
elif im.ndim == 4:
|
|
if force_rgba2rgb:
|
|
im = im[:, :, :3]
|
|
else:
|
|
raise ValueError(f"{str(path)} has {im.ndim} channels!")
|
|
else:
|
|
if chn.lower() == 'rgb':
|
|
im = bgr2rgb(im)
|
|
elif chn.lower() == 'bgr':
|
|
pass
|
|
|
|
if dtype == 'float32':
|
|
im = im.astype(np.float32) / 255.
|
|
elif dtype == 'float64':
|
|
im = im.astype(np.float64) / 255.
|
|
elif dtype == 'uint8':
|
|
pass
|
|
else:
|
|
sys.exit('Please input corrected dtype: float32, float64 or uint8!')
|
|
|
|
return im
|
|
|
|
def imwrite(im_in, path, chn='rgb', dtype_in='float32', qf=None):
|
|
'''
|
|
Save image.
|
|
Input:
|
|
im: h x w x c, numpy tensor
|
|
path: the saving path
|
|
chn: the channel order of the im,
|
|
'''
|
|
im = im_in.copy()
|
|
if isinstance(path, str):
|
|
path = Path(path)
|
|
if dtype_in != 'uint8':
|
|
im = img_as_ubyte(im)
|
|
|
|
if chn.lower() == 'rgb' and im.ndim == 3:
|
|
im = rgb2bgr(im)
|
|
|
|
if qf is not None and path.suffix.lower() in ['.jpg', '.jpeg']:
|
|
flag = cv2.imwrite(str(path), im, [int(cv2.IMWRITE_JPEG_QUALITY), int(qf)])
|
|
else:
|
|
flag = cv2.imwrite(str(path), im)
|
|
|
|
return flag
|
|
|
|
def jpeg_compress(im, qf, chn_in='rgb'):
|
|
'''
|
|
Input:
|
|
im: h x w x 3 array
|
|
qf: compress factor, (0, 100]
|
|
chn_in: 'rgb' or 'bgr'
|
|
Return:
|
|
Compressed Image with channel order: chn_in
|
|
'''
|
|
# transform to BGR channle and uint8 data type
|
|
im_bgr = rgb2bgr(im) if chn_in.lower() == 'rgb' else im
|
|
if im.dtype != np.dtype('uint8'): im_bgr = img_as_ubyte(im_bgr)
|
|
|
|
# JPEG compress
|
|
flag, encimg = cv2.imencode('.jpg', im_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), qf])
|
|
assert flag
|
|
im_jpg_bgr = cv2.imdecode(encimg, 1) # uint8, BGR
|
|
|
|
# transform back to original channel and the original data type
|
|
im_out = bgr2rgb(im_jpg_bgr) if chn_in.lower() == 'rgb' else im_jpg_bgr
|
|
if im.dtype != np.dtype('uint8'): im_out = img_as_float32(im_out).astype(im.dtype)
|
|
return im_out
|
|
|
|
# ------------------------Augmentation-----------------------------
|
|
def data_aug_np(image, mode):
|
|
'''
|
|
Performs data augmentation of the input image
|
|
Input:
|
|
image: a cv2 (OpenCV) image
|
|
mode: int. Choice of transformation to apply to the image
|
|
0 - no transformation
|
|
1 - flip up and down
|
|
2 - rotate counterwise 90 degree
|
|
3 - rotate 90 degree and flip up and down
|
|
4 - rotate 180 degree
|
|
5 - rotate 180 degree and flip
|
|
6 - rotate 270 degree
|
|
7 - rotate 270 degree and flip
|
|
'''
|
|
if mode == 0:
|
|
# original
|
|
out = image
|
|
elif mode == 1:
|
|
# flip up and down
|
|
out = np.flipud(image)
|
|
elif mode == 2:
|
|
# rotate counterwise 90 degree
|
|
out = np.rot90(image)
|
|
elif mode == 3:
|
|
# rotate 90 degree and flip up and down
|
|
out = np.rot90(image)
|
|
out = np.flipud(out)
|
|
elif mode == 4:
|
|
# rotate 180 degree
|
|
out = np.rot90(image, k=2)
|
|
elif mode == 5:
|
|
# rotate 180 degree and flip
|
|
out = np.rot90(image, k=2)
|
|
out = np.flipud(out)
|
|
elif mode == 6:
|
|
# rotate 270 degree
|
|
out = np.rot90(image, k=3)
|
|
elif mode == 7:
|
|
# rotate 270 degree and flip
|
|
out = np.rot90(image, k=3)
|
|
out = np.flipud(out)
|
|
else:
|
|
raise Exception('Invalid choice of image transformation')
|
|
|
|
return out.copy()
|
|
|
|
def inverse_data_aug_np(image, mode):
|
|
'''
|
|
Performs inverse data augmentation of the input image
|
|
'''
|
|
if mode == 0:
|
|
# original
|
|
out = image
|
|
elif mode == 1:
|
|
out = np.flipud(image)
|
|
elif mode == 2:
|
|
out = np.rot90(image, axes=(1,0))
|
|
elif mode == 3:
|
|
out = np.flipud(image)
|
|
out = np.rot90(out, axes=(1,0))
|
|
elif mode == 4:
|
|
out = np.rot90(image, k=2, axes=(1,0))
|
|
elif mode == 5:
|
|
out = np.flipud(image)
|
|
out = np.rot90(out, k=2, axes=(1,0))
|
|
elif mode == 6:
|
|
out = np.rot90(image, k=3, axes=(1,0))
|
|
elif mode == 7:
|
|
# rotate 270 degree and flip
|
|
out = np.flipud(image)
|
|
out = np.rot90(out, k=3, axes=(1,0))
|
|
else:
|
|
raise Exception('Invalid choice of image transformation')
|
|
|
|
return out
|
|
|
|
# ----------------------Visualization----------------------------
|
|
def imshow(x, title=None, cbar=False):
|
|
import matplotlib.pyplot as plt
|
|
plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
|
|
if title:
|
|
plt.title(title)
|
|
if cbar:
|
|
plt.colorbar()
|
|
plt.show()
|
|
|
|
def imblend_with_mask(im, mask, alpha=0.25):
|
|
"""
|
|
Input:
|
|
im, mask: h x w x c numpy array, uint8, [0, 255]
|
|
alpha: scaler in [0.0, 1.0]
|
|
"""
|
|
edge_map = cv2.Canny(mask, 100, 200).astype(np.float32)[:, :, None] / 255.
|
|
|
|
assert mask.dtype == np.uint8
|
|
mask = mask.astype(np.float32) / 255.
|
|
if mask.ndim == 2:
|
|
mask = mask[:, :, None]
|
|
|
|
back_color = np.array([159, 121, 238], dtype=np.float32).reshape((1,1,3))
|
|
blend = im.astype(np.float32) * alpha + (1 - alpha) * back_color
|
|
blend = np.clip(blend, 0, 255)
|
|
out = im.astype(np.float32) * (1 - mask) + blend * mask
|
|
|
|
# paste edge
|
|
out = out * (1 - edge_map) + np.array([0,255,0], dtype=np.float32).reshape((1,1,3)) * edge_map
|
|
|
|
return out.astype(np.uint8)
|
|
# -----------------------Covolution------------------------------
|
|
def imgrad(im, pading_mode='mirror'):
|
|
'''
|
|
Calculate image gradient.
|
|
Input:
|
|
im: h x w x c numpy array
|
|
'''
|
|
from scipy.ndimage import correlate # lazy import
|
|
wx = np.array([[0, 0, 0],
|
|
[-1, 1, 0],
|
|
[0, 0, 0]], dtype=np.float32)
|
|
wy = np.array([[0, -1, 0],
|
|
[0, 1, 0],
|
|
[0, 0, 0]], dtype=np.float32)
|
|
if im.ndim == 3:
|
|
gradx = np.stack(
|
|
[correlate(im[:,:,c], wx, mode=pading_mode) for c in range(im.shape[2])],
|
|
axis=2
|
|
)
|
|
grady = np.stack(
|
|
[correlate(im[:,:,c], wy, mode=pading_mode) for c in range(im.shape[2])],
|
|
axis=2
|
|
)
|
|
grad = np.concatenate((gradx, grady), axis=2)
|
|
else:
|
|
gradx = correlate(im, wx, mode=pading_mode)
|
|
grady = correlate(im, wy, mode=pading_mode)
|
|
grad = np.stack((gradx, grady), axis=2)
|
|
|
|
return {'gradx': gradx, 'grady': grady, 'grad':grad}
|
|
|
|
def imgrad_fft(im):
|
|
'''
|
|
Calculate image gradient.
|
|
Input:
|
|
im: h x w x c numpy array
|
|
'''
|
|
wx = np.rot90(np.array([[0, 0, 0],
|
|
[-1, 1, 0],
|
|
[0, 0, 0]], dtype=np.float32), k=2)
|
|
gradx = convfft(im, wx)
|
|
wy = np.rot90(np.array([[0, -1, 0],
|
|
[0, 1, 0],
|
|
[0, 0, 0]], dtype=np.float32), k=2)
|
|
grady = convfft(im, wy)
|
|
grad = np.concatenate((gradx, grady), axis=2)
|
|
|
|
return {'gradx': gradx, 'grady': grady, 'grad':grad}
|
|
|
|
def convfft(im, weight):
|
|
'''
|
|
Convolution with FFT
|
|
Input:
|
|
im: h1 x w1 x c numpy array
|
|
weight: h2 x w2 numpy array
|
|
Output:
|
|
out: h1 x w1 x c numpy array
|
|
'''
|
|
axes = (0,1)
|
|
otf = psf2otf(weight, im.shape[:2])
|
|
if im.ndim == 3:
|
|
otf = np.tile(otf[:, :, None], (1,1,im.shape[2]))
|
|
out = fft.ifft2(fft.fft2(im, axes=axes) * otf, axes=axes).real
|
|
return out
|
|
|
|
def psf2otf(psf, shape):
|
|
"""
|
|
MATLAB psf2otf function.
|
|
Borrowed from https://github.com/aboucaud/pypher/blob/master/pypher/pypher.py.
|
|
Input:
|
|
psf : h x w numpy array
|
|
shape : list or tuple, output shape of the OTF array
|
|
Output:
|
|
otf : OTF array with the desirable shape
|
|
"""
|
|
if np.all(psf == 0):
|
|
return np.zeros_like(psf)
|
|
|
|
inshape = psf.shape
|
|
# Pad the PSF to outsize
|
|
psf = zero_pad(psf, shape, position='corner')
|
|
|
|
# Circularly shift OTF so that the 'center' of the PSF is [0,0] element of the array
|
|
for axis, axis_size in enumerate(inshape):
|
|
psf = np.roll(psf, -int(axis_size / 2), axis=axis)
|
|
|
|
# Compute the OTF
|
|
otf = fft.fft2(psf)
|
|
|
|
# Estimate the rough number of operations involved in the FFT
|
|
# and discard the PSF imaginary part if within roundoff error
|
|
# roundoff error = machine epsilon = sys.float_info.epsilon
|
|
# or np.finfo().eps
|
|
n_ops = np.sum(psf.size * np.log2(psf.shape))
|
|
otf = np.real_if_close(otf, tol=n_ops)
|
|
|
|
return otf
|
|
|
|
def convtorch(im, weight, mode='reflect'):
|
|
'''
|
|
Image convolution with pytorch
|
|
Input:
|
|
im: b x c_in x h x w torch tensor
|
|
weight: c_out x c_in x k x k torch tensor
|
|
Output:
|
|
out: c x h x w torch tensor
|
|
'''
|
|
radius = weight.shape[-1]
|
|
chn = im.shape[1]
|
|
im_pad = torch.nn.functional.pad(im, pad=(radius // 2, )*4, mode=mode)
|
|
out = torch.nn.functional.conv2d(im_pad, weight, padding=0, groups=chn)
|
|
return out
|
|
|
|
# ----------------------Patch Cropping----------------------------
|
|
def random_crop(im, pch_size):
|
|
'''
|
|
Randomly crop a patch from the give image.
|
|
'''
|
|
h, w = im.shape[:2]
|
|
# padding if necessary
|
|
if h < pch_size or w < pch_size:
|
|
pad_h = min(max(0, pch_size - h), h)
|
|
pad_w = min(max(0, pch_size - w), w)
|
|
im = cv2.copyMakeBorder(im, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
|
|
|
|
h, w = im.shape[:2]
|
|
if h == pch_size:
|
|
ind_h = 0
|
|
elif h > pch_size:
|
|
ind_h = random.randint(0, h-pch_size)
|
|
else:
|
|
raise ValueError('Image height is smaller than the patch size')
|
|
if w == pch_size:
|
|
ind_w = 0
|
|
elif w > pch_size:
|
|
ind_w = random.randint(0, w-pch_size)
|
|
else:
|
|
raise ValueError('Image width is smaller than the patch size')
|
|
|
|
im_pch = im[ind_h:ind_h+pch_size, ind_w:ind_w+pch_size,]
|
|
|
|
return im_pch
|
|
|
|
class ToTensor:
|
|
def __init__(self, max_value=1.0):
|
|
self.max_value = max_value
|
|
|
|
def __call__(self, im):
|
|
assert isinstance(im, np.ndarray)
|
|
if im.ndim == 2:
|
|
im = im[:, :, np.newaxis]
|
|
if im.dtype == np.uint8:
|
|
assert self.max_value == 255.
|
|
out = torch.from_numpy(im.astype(np.float32).transpose(2,0,1) / self.max_value)
|
|
else:
|
|
assert self.max_value == 1.0
|
|
out = torch.from_numpy(im.transpose(2,0,1))
|
|
return out
|
|
|
|
class RandomCrop:
|
|
def __init__(self, pch_size, pass_crop=False):
|
|
self.pch_size = pch_size
|
|
self.pass_crop = pass_crop
|
|
|
|
def __call__(self, im):
|
|
if self.pass_crop:
|
|
return im
|
|
if isinstance(im, list) or isinstance(im, tuple):
|
|
out = []
|
|
for current_im in im:
|
|
out.append(random_crop(current_im, self.pch_size))
|
|
else:
|
|
out = random_crop(im, self.pch_size)
|
|
return out
|
|
|
|
class ImageSpliterNp:
|
|
def __init__(self, im, pch_size, stride, sf=1):
|
|
'''
|
|
Input:
|
|
im: h x w x c, numpy array, [0, 1], low-resolution image in SR
|
|
pch_size, stride: patch setting
|
|
sf: scale factor in image super-resolution
|
|
'''
|
|
assert stride <= pch_size
|
|
self.stride = stride
|
|
self.pch_size = pch_size
|
|
self.sf = sf
|
|
|
|
if im.ndim == 2:
|
|
im = im[:, :, None]
|
|
|
|
height, width, chn = im.shape
|
|
self.height_starts_list = self.extract_starts(height)
|
|
self.width_starts_list = self.extract_starts(width)
|
|
self.length = self.__len__()
|
|
self.num_pchs = 0
|
|
|
|
self.im_ori = im
|
|
self.im_res = np.zeros([height*sf, width*sf, chn], dtype=im.dtype)
|
|
self.pixel_count = np.zeros([height*sf, width*sf, chn], dtype=im.dtype)
|
|
|
|
def extract_starts(self, length):
|
|
starts = list(range(0, length, self.stride))
|
|
if starts[-1] + self.pch_size > length:
|
|
starts[-1] = length - self.pch_size
|
|
return starts
|
|
|
|
def __len__(self):
|
|
return len(self.height_starts_list) * len(self.width_starts_list)
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
if self.num_pchs < self.length:
|
|
w_start_idx = self.num_pchs // len(self.height_starts_list)
|
|
w_start = self.width_starts_list[w_start_idx] * self.sf
|
|
w_end = w_start + self.pch_size * self.sf
|
|
|
|
h_start_idx = self.num_pchs % len(self.height_starts_list)
|
|
h_start = self.height_starts_list[h_start_idx] * self.sf
|
|
h_end = h_start + self.pch_size * self.sf
|
|
|
|
pch = self.im_ori[h_start:h_end, w_start:w_end,]
|
|
self.w_start, self.w_end = w_start, w_end
|
|
self.h_start, self.h_end = h_start, h_end
|
|
|
|
self.num_pchs += 1
|
|
else:
|
|
raise StopIteration(0)
|
|
|
|
return pch, (h_start, h_end, w_start, w_end)
|
|
|
|
def update(self, pch_res, index_infos):
|
|
'''
|
|
Input:
|
|
pch_res: pch_size x pch_size x 3, [0,1]
|
|
index_infos: (h_start, h_end, w_start, w_end)
|
|
'''
|
|
if index_infos is None:
|
|
w_start, w_end = self.w_start, self.w_end
|
|
h_start, h_end = self.h_start, self.h_end
|
|
else:
|
|
h_start, h_end, w_start, w_end = index_infos
|
|
|
|
self.im_res[h_start:h_end, w_start:w_end] += pch_res
|
|
self.pixel_count[h_start:h_end, w_start:w_end] += 1
|
|
|
|
def gather(self):
|
|
assert np.all(self.pixel_count != 0)
|
|
return self.im_res / self.pixel_count
|
|
|
|
class ImageSpliterTh:
|
|
def __init__(self, im, pch_size, stride, sf=1, extra_bs=1, weight_type='Gaussian'):
|
|
'''
|
|
Input:
|
|
im: n x c x h x w, torch tensor, float, low-resolution image in SR
|
|
pch_size, stride: patch setting
|
|
sf: scale factor in image super-resolution
|
|
pch_bs: aggregate pchs to processing, only used when inputing single image
|
|
'''
|
|
assert weight_type in ['Gaussian', 'ones']
|
|
self.weight_type = weight_type
|
|
assert stride <= pch_size
|
|
self.stride = stride
|
|
self.pch_size = pch_size
|
|
self.sf = sf
|
|
self.extra_bs = extra_bs
|
|
|
|
bs, chn, height, width= im.shape
|
|
self.true_bs = bs
|
|
|
|
self.height_starts_list = self.extract_starts(height)
|
|
self.width_starts_list = self.extract_starts(width)
|
|
self.starts_list = []
|
|
for ii in self.height_starts_list:
|
|
for jj in self.width_starts_list:
|
|
self.starts_list.append([ii, jj])
|
|
|
|
self.length = self.__len__()
|
|
self.count_pchs = 0
|
|
|
|
self.im_ori = im
|
|
self.dtype = torch.float64
|
|
self.im_res = torch.zeros([bs, chn, height*sf, width*sf], dtype=self.dtype, device=im.device)
|
|
self.pixel_count = torch.zeros([bs, chn, height*sf, width*sf], dtype=self.dtype, device=im.device)
|
|
|
|
def extract_starts(self, length):
|
|
if length <= self.pch_size:
|
|
starts = [0,]
|
|
else:
|
|
starts = list(range(0, length, self.stride))
|
|
for ii in range(len(starts)):
|
|
if starts[ii] + self.pch_size > length:
|
|
starts[ii] = length - self.pch_size
|
|
starts = sorted(set(starts), key=starts.index)
|
|
return starts
|
|
|
|
def __len__(self):
|
|
return len(self.height_starts_list) * len(self.width_starts_list)
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
if self.count_pchs < self.length:
|
|
index_infos = []
|
|
current_starts_list = self.starts_list[self.count_pchs:self.count_pchs+self.extra_bs]
|
|
for ii, (h_start, w_start) in enumerate(current_starts_list):
|
|
w_end = w_start + self.pch_size
|
|
h_end = h_start + self.pch_size
|
|
current_pch = self.im_ori[:, :, h_start:h_end, w_start:w_end]
|
|
if ii == 0:
|
|
pch = current_pch
|
|
else:
|
|
pch = torch.cat([pch, current_pch], dim=0)
|
|
|
|
h_start *= self.sf
|
|
h_end *= self.sf
|
|
w_start *= self.sf
|
|
w_end *= self.sf
|
|
index_infos.append([h_start, h_end, w_start, w_end])
|
|
|
|
self.count_pchs += len(current_starts_list)
|
|
else:
|
|
raise StopIteration()
|
|
|
|
return pch, index_infos
|
|
|
|
def update(self, pch_res, index_infos):
|
|
'''
|
|
Input:
|
|
pch_res: (n*extra_bs) x c x pch_size x pch_size, float
|
|
index_infos: [(h_start, h_end, w_start, w_end),]
|
|
'''
|
|
assert pch_res.shape[0] % self.true_bs == 0
|
|
pch_list = torch.split(pch_res, self.true_bs, dim=0)
|
|
assert len(pch_list) == len(index_infos)
|
|
for ii, (h_start, h_end, w_start, w_end) in enumerate(index_infos):
|
|
current_pch = pch_list[ii].type(self.dtype)
|
|
current_weight = self.get_weight(current_pch.shape[-2], current_pch.shape[-1])
|
|
self.im_res[:, :, h_start:h_end, w_start:w_end] += current_pch * current_weight
|
|
self.pixel_count[:, :, h_start:h_end, w_start:w_end] += current_weight
|
|
|
|
@staticmethod
|
|
def generate_kernel_1d(ksize):
|
|
sigma = 0.3 * ((ksize - 1) * 0.5 - 1) + 0.8 # opencv default setting
|
|
if ksize % 2 == 0:
|
|
kernel = cv2.getGaussianKernel(ksize=ksize+1, sigma=sigma, ktype=cv2.CV_64F)
|
|
kernel = kernel[1:, ]
|
|
else:
|
|
kernel = cv2.getGaussianKernel(ksize=ksize, sigma=sigma, ktype=cv2.CV_64F)
|
|
|
|
return kernel
|
|
|
|
def get_weight(self, height, width):
|
|
if self.weight_type == 'ones':
|
|
kernel = torch.ones(1, 1, height, width)
|
|
elif self.weight_type == 'Gaussian':
|
|
kernel_h = self.generate_kernel_1d(height).reshape(-1, 1)
|
|
kernel_w = self.generate_kernel_1d(width).reshape(1, -1)
|
|
kernel = np.matmul(kernel_h, kernel_w)
|
|
kernel = torch.from_numpy(kernel).unsqueeze(0).unsqueeze(0) # 1 x 1 x pch_size x pch_size
|
|
else:
|
|
raise ValueError(f"Unsupported weight type: {self.weight_type}")
|
|
|
|
return kernel.to(dtype=self.dtype, device=self.im_ori.device)
|
|
|
|
def gather(self):
|
|
assert torch.all(self.pixel_count != 0)
|
|
return self.im_res.div(self.pixel_count)
|
|
|
|
# ----------------------Patch Cliping----------------------------
|
|
class Clamper:
|
|
def __init__(self, min_max=(-1, 1)):
|
|
self.min_bound, self.max_bound = min_max[0], min_max[1]
|
|
|
|
def __call__(self, im):
|
|
if isinstance(im, np.ndarray):
|
|
return np.clip(im, a_min=self.min_bound, a_max=self.max_bound)
|
|
elif isinstance(im, torch.Tensor):
|
|
return torch.clamp(im, min=self.min_bound, max=self.max_bound)
|
|
else:
|
|
raise TypeError(f'ndarray or Tensor expected, got {type(im)}')
|
|
|
|
# ----------------------Interpolation----------------------------
|
|
class Bicubic:
|
|
def __init__(self, scale=None, out_shape=None, activate_matlab=True, resize_back=False):
|
|
self.scale = scale
|
|
self.activate_matlab = activate_matlab
|
|
self.out_shape = out_shape
|
|
self.resize_back = resize_back
|
|
|
|
def __call__(self, im):
|
|
if self.activate_matlab:
|
|
out = imresize_np(im, scale=self.scale)
|
|
if self.resize_back:
|
|
out = imresize_np(out, scale=1/self.scale)
|
|
else:
|
|
out = cv2.resize(
|
|
im,
|
|
dsize=self.out_shape,
|
|
fx=self.scale,
|
|
fy=self.scale,
|
|
interpolation=cv2.INTER_CUBIC,
|
|
)
|
|
if self.resize_back:
|
|
out = cv2.resize(
|
|
out,
|
|
dsize=self.out_shape,
|
|
fx=1/self.scale,
|
|
fy=1/self.scale,
|
|
interpolation=cv2.INTER_CUBIC,
|
|
)
|
|
return out
|
|
|
|
class SmallestMaxSize:
|
|
def __init__(self, max_size, pass_resize=False, interpolation=None):
|
|
self.pass_resize = pass_resize
|
|
self.max_size = max_size
|
|
self.interpolation = interpolation
|
|
self.str2mode = {
|
|
'nearest': cv2.INTER_NEAREST_EXACT,
|
|
'bilinear': cv2.INTER_LINEAR,
|
|
'bicubic': cv2.INTER_CUBIC
|
|
}
|
|
if self.interpolation is not None:
|
|
assert interpolation in self.str2mode, f"Not supported interpolation mode: {interpolation}"
|
|
|
|
def get_interpolation(self, size):
|
|
if self.interpolation is None:
|
|
if size < self.max_size: # upsampling
|
|
interpolation = cv2.INTER_CUBIC
|
|
else: # downsampling
|
|
interpolation = cv2.INTER_AREA
|
|
else:
|
|
interpolation = self.str2mode[self.interpolation]
|
|
|
|
return interpolation
|
|
|
|
def __call__(self, im):
|
|
h, w = im.shape[:2]
|
|
if self.pass_resize or min(h, w) == self.max_size:
|
|
out = im
|
|
else:
|
|
if h < w:
|
|
dsize = (int(self.max_size * w / h), self.max_size)
|
|
out = cv2.resize(im, dsize=dsize, interpolation=self.get_interpolation(h))
|
|
else:
|
|
dsize = (self.max_size, int(self.max_size * h / w))
|
|
out = cv2.resize(im, dsize=dsize, interpolation=self.get_interpolation(w))
|
|
if out.dtype == np.uint8:
|
|
out = np.clip(out, 0, 255)
|
|
else:
|
|
out = np.clip(out, 0, 1.0)
|
|
|
|
return out
|
|
|
|
# ----------------------augmentation----------------------------
|
|
class SpatialAug:
|
|
def __init__(self, pass_aug, only_hflip=False, only_vflip=False, only_hvflip=False):
|
|
self.only_hflip = only_hflip
|
|
self.only_vflip = only_vflip
|
|
self.only_hvflip = only_hvflip
|
|
self.pass_aug = pass_aug
|
|
|
|
def __call__(self, im, flag=None):
|
|
if self.pass_aug:
|
|
return im
|
|
|
|
if flag is None:
|
|
if self.only_hflip:
|
|
flag = random.choice([0, 5])
|
|
elif self.only_vflip:
|
|
flag = random.choice([0, 1])
|
|
elif self.only_hvflip:
|
|
flag = random.choice([0, 1, 5])
|
|
else:
|
|
flag = random.randint(0, 7)
|
|
|
|
if isinstance(im, list) or isinstance(im, tuple):
|
|
out = []
|
|
for current_im in im:
|
|
out.append(data_aug_np(current_im, flag))
|
|
else:
|
|
out = data_aug_np(im, flag)
|
|
return out
|
|
|
|
if __name__ == '__main__':
|
|
im = np.random.randn(64, 64, 3).astype(np.float32)
|
|
|
|
grad1 = imgrad(im)['grad']
|
|
grad2 = imgrad_fft(im)['grad']
|
|
|
|
error = np.abs(grad1 -grad2).max()
|
|
mean_error = np.abs(grad1 -grad2).mean()
|
|
print('The largest error is {:.2e}'.format(error))
|
|
print('The mean error is {:.2e}'.format(mean_error))
|