#!/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))