plt.figure(figsize=figsize)

plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')

if title:

plt.title(title)

if cbar:

plt.colorbar()

plt.figure(figsize=figsize)

ax3 = plt.axes(projection='3d')

w, h = Z.shape[:2]

xx = np.arange(0,w,1)

yy = np.arange(0,h,1)

X, Y = np.meshgrid(xx, yy)

ax3.plot_surface(X,Y,Z,cmap=cmap)

#ax3.contour(X,Y,Z, zdim='z',offset=-2，cmap=cmap)

paths = None # return None if dataroot is None

if dataroot is not None:

paths = sorted(_get_paths_from_images(dataroot))

assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)

images = []

for dirpath, _, fnames in sorted(os.walk(path)):

for fname in sorted(fnames):

if is_image_file(fname):

img_path = os.path.join(dirpath, fname)

images.append(img_path)

assert images, '{:s} has no valid image file'.format(path)

w, h = img.shape[:2]

patches = []

if w > p_max and h > p_max:

w1 = list(np.arange(0, w-p_size, p_size-p_overlap, dtype=np.int))

h1 = list(np.arange(0, h-p_size, p_size-p_overlap, dtype=np.int))

w1.append(w-p_size)

h1.append(h-p_size)

print(w1)

print(h1)

for i in w1:

for j in h1:

patches.append(img[i:i+p_size, j:j+p_size,:])

else:

patches.append(img)

"""

imgs: list, N images of size WxHxC

"""

img_name, ext = os.path.splitext(os.path.basename(img_path))

for i, img in enumerate(imgs):

if img.ndim == 3:

img = img[:, :, [2, 1, 0]]

new_path = os.path.dirname(img_path)+img_name+str('_{:04d}'.format(i))+ext

"""

split the large images from original_dataroot into small overlapped images with size (p_size)x(p_size),

and save them into taget_dataroot; only the images with larger size than (p_max)x(p_max)

will be splitted.

Args:

original_dataroot:

taget_dataroot:

p_size: size of small images

p_overlap: patch size in training is a good choice

p_max: images with smaller size than (p_max)x(p_max) keep unchanged.

"""

paths = get_image_paths(original_dataroot)

for img_path in paths:

# img_name, ext = os.path.splitext(os.path.basename(img_path))

img = imread_uint(img_path, n_channels=n_channels)

patches = patches_from_image(img, p_size, p_overlap, p_max)

imssave(patches, taget_dataroot+os.path.basename(img_path))

#if original_dataroot == taget_dataroot:

if isinstance(paths, str):

mkdir(paths)

else:

for path in paths:

if os.path.exists(path):

new_name = path + '_archived_' + get_timestamp()

print('Path already exists. Rename it to [{:s}]'.format(new_name))

os.rename(path, new_name)

# input: path

# output: HxWx3(RGB or GGG), or HxWx1 (G)

if n_channels == 1:

img = cv2.imread(path, 0) # cv2.IMREAD_GRAYSCALE

img = np.expand_dims(img, axis=2) # HxWx1

elif n_channels == 3:

img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGR or G

if img.ndim == 2:

img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # GGG

else:

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB

img = np.squeeze(img)

if img.ndim == 3:

img = img[:, :, [2, 1, 0]]

img = np.squeeze(img)

if img.ndim == 3:

img = img[:, :, [2, 1, 0]]

# read image by cv2

# return: Numpy float32, HWC, BGR, [0,1]

img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # cv2.IMREAD_GRAYSCALE

img = img.astype(np.float32) / 255.

if img.ndim == 2:

img = np.expand_dims(img, axis=2)

# some images have 4 channels

if img.shape[2] > 3:

img = img[:, :, :3]

img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()

if img.ndim == 3:

img = np.transpose(img, (1, 2, 0))

img = img.data.squeeze().float().cpu().numpy()

if img.ndim == 3:

img = np.transpose(img, (1, 2, 0))

img = img.data.squeeze().float().cpu().numpy()

if img.ndim == 3:

img = np.transpose(img, (1, 2, 0))

elif img.ndim == 2:

img = np.expand_dims(img, axis=2)

'''

Converts a torch Tensor into an image Numpy array of BGR channel order

Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order

Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)

'''

tensor = tensor.squeeze().float().cpu().clamp_(*min_max) # squeeze first, then clamp

tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) # to range [0,1]

n_dim = tensor.dim()

if n_dim == 4:

n_img = len(tensor)

img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()

img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR

elif n_dim == 3:

img_np = tensor.numpy()

img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR

elif n_dim == 2:

img_np = tensor.numpy()

else:

raise TypeError(

'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))

if out_type == np.uint8:

img_np = (img_np * 255.0).round()

# Important. Unlike matlab, numpy.unit8() WILL NOT round by default.

'''Kai Zhang (github: https://github.com/cszn)

'''

if mode == 0:

return img

elif mode == 1:

return np.flipud(np.rot90(img))

elif mode == 2:

return np.flipud(img)

elif mode == 3:

return np.rot90(img, k=3)

elif mode == 4:

return np.flipud(np.rot90(img, k=2))

elif mode == 5:

return np.rot90(img)

elif mode == 6:

return np.rot90(img, k=2)

elif mode == 7:

'''Kai Zhang (github: https://github.com/cszn)

'''

if mode == 0:

return img

elif mode == 1:

return img.rot90(1, [2, 3]).flip([2])

elif mode == 2:

return img.flip([2])

elif mode == 3:

return img.rot90(3, [2, 3])

elif mode == 4:

return img.rot90(2, [2, 3]).flip([2])

elif mode == 5:

return img.rot90(1, [2, 3])

elif mode == 6:

return img.rot90(2, [2, 3])

elif mode == 7:

'''Kai Zhang (github: https://github.com/cszn)

'''

img_size = img.size()

img_np = img.data.cpu().numpy()

if len(img_size) == 3:

img_np = np.transpose(img_np, (1, 2, 0))

elif len(img_size) == 4:

img_np = np.transpose(img_np, (2, 3, 1, 0))

img_np = augment_img(img_np, mode=mode)

img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))

if len(img_size) == 3:

img_tensor = img_tensor.permute(2, 0, 1)

elif len(img_size) == 4:

img_tensor = img_tensor.permute(3, 2, 0, 1)

if mode == 0:

return img

elif mode == 1:

return img.transpose(1, 0, 2)

elif mode == 2:

return img[::-1, :, :]

elif mode == 3:

img = img[::-1, :, :]

img = img.transpose(1, 0, 2)

return img

elif mode == 4:

return img[:, ::-1, :]

elif mode == 5:

img = img[:, ::-1, :]

img = img.transpose(1, 0, 2)

return img

elif mode == 6:

img = img[:, ::-1, :]

img = img[::-1, :, :]

return img

elif mode == 7:

img = img[:, ::-1, :]

img = img[::-1, :, :]

img = img.transpose(1, 0, 2)

# horizontal flip OR rotate

hflip = hflip and random.random() < 0.5

vflip = rot and random.random() < 0.5

rot90 = rot and random.random() < 0.5

def _augment(img):

if hflip:

img = img[:, ::-1, :]

if vflip:

img = img[::-1, :, :]

if rot90:

img = img.transpose(1, 0, 2)

return img

# img_in: Numpy, HWC or HW

img = np.copy(img_in)

if img.ndim == 2:

H, W = img.shape

H_r, W_r = H % scale, W % scale

img = img[:H - H_r, :W - W_r]

elif img.ndim == 3:

H, W, C = img.shape

H_r, W_r = H % scale, W % scale

img = img[:H - H_r, :W - W_r, :]

else:

raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))

# img_in: Numpy, HWC or HW

img = np.copy(img_in)

h, w = img.shape[:2]

img = img[border:h-border, border:w-border]

'''same as matlab rgb2ycbcr

only_y: only return Y channel

Input:

uint8, [0, 255]

float, [0, 1]

'''

in_img_type = img.dtype

img.astype(np.float32)

if in_img_type != np.uint8:

img *= 255.

# convert

if only_y:

rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0

else:

rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],

[24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]

if in_img_type == np.uint8:

rlt = rlt.round()

else:

rlt /= 255.

'''same as matlab ycbcr2rgb

Input:

uint8, [0, 255]

float, [0, 1]

'''

in_img_type = img.dtype

img.astype(np.float32)

if in_img_type != np.uint8:

img *= 255.

# convert

rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],

[0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]

if in_img_type == np.uint8:

rlt = rlt.round()

else:

rlt /= 255.

'''bgr version of rgb2ycbcr

only_y: only return Y channel

Input:

uint8, [0, 255]

float, [0, 1]

'''

in_img_type = img.dtype

img.astype(np.float32)

if in_img_type != np.uint8:

img *= 255.

# convert

if only_y:

rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0

else:

rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],

[65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]

if in_img_type == np.uint8:

rlt = rlt.round()

else:

rlt /= 255.

# conversion among BGR, gray and y

if in_c == 3 and tar_type == 'gray': # BGR to gray

gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]

return [np.expand_dims(img, axis=2) for img in gray_list]

elif in_c == 3 and tar_type == 'y': # BGR to y

y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]

return [np.expand_dims(img, axis=2) for img in y_list]

elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR

return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]

else:

# img1 and img2 have range [0, 255]

#img1 = img1.squeeze()

#img2 = img2.squeeze()

if not img1.shape == img2.shape:

raise ValueError('Input images must have the same dimensions.')

h, w = img1.shape[:2]

img1 = img1[border:h-border, border:w-border]

img2 = img2[border:h-border, border:w-border]

img1 = img1.astype(np.float64)

img2 = img2.astype(np.float64)

mse = np.mean((img1 - img2)**2)

if mse == 0:

return float('inf')

'''calculate SSIM

the same outputs as MATLAB's

img1, img2: [0, 255]

'''

#img1 = img1.squeeze()

#img2 = img2.squeeze()

if not img1.shape == img2.shape:

raise ValueError('Input images must have the same dimensions.')

h, w = img1.shape[:2]

img1 = img1[border:h-border, border:w-border]

img2 = img2[border:h-border, border:w-border]

if img1.ndim == 2:

return ssim(img1, img2)

elif img1.ndim == 3:

if img1.shape[2] == 3:

ssims = []

for i in range(3):

ssims.append(ssim(img1[:,:,i], img2[:,:,i]))

return np.array(ssims).mean()

elif img1.shape[2] == 1:

return ssim(np.squeeze(img1), np.squeeze(img2))

else:

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))

absx = torch.abs(x)

absx2 = absx**2

absx3 = absx**3

return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \

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

# Now the scale should be the same for H and W

# input: img: pytorch tensor, CHW or HW [0,1]

# output: CHW or HW [0,1] w/o round

need_squeeze = True if img.dim() == 2 else False

if need_squeeze:

img.unsqueeze_(0)

in_C, in_H, in_W = 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_C, in_H + sym_len_Hs + sym_len_He, in_W)

img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)

sym_patch = img[:, :sym_len_Hs, :]

inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()

sym_patch_inv = sym_patch.index_select(1, inv_idx)

img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)

sym_patch = img[:, -sym_len_He:, :]

inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()

sym_patch_inv = sym_patch.index_select(1, inv_idx)

img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)

out_1 = torch.FloatTensor(in_C, out_H, in_W)

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[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])

# process W dimension

# symmetric copying

out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)

out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)

sym_patch = out_1[:, :, :sym_len_Ws]

inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()

sym_patch_inv = sym_patch.index_select(2, inv_idx)

out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)

sym_patch = out_1[:, :, -sym_len_We:]

inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()

sym_patch_inv = sym_patch.index_select(2, inv_idx)

out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)

out_2 = torch.FloatTensor(in_C, out_H, out_W)

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[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i])

if need_squeeze:

out_2.squeeze_()

# 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_()