mirror of
https://github.com/Stability-AI/stablediffusion.git
synced 2024-12-22 23:55:00 +00:00
916 lines
No EOL
28 KiB
Python
916 lines
No EOL
28 KiB
Python
import os
|
||
import math
|
||
import random
|
||
import numpy as np
|
||
import torch
|
||
import cv2
|
||
from torchvision.utils import make_grid
|
||
from datetime import datetime
|
||
#import matplotlib.pyplot as plt # TODO: check with Dominik, also bsrgan.py vs bsrgan_light.py
|
||
|
||
|
||
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# Kai Zhang (github: https://github.com/cszn)
|
||
# 03/Mar/2019
|
||
# --------------------------------------------
|
||
# https://github.com/twhui/SRGAN-pyTorch
|
||
# https://github.com/xinntao/BasicSR
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif']
|
||
|
||
|
||
def is_image_file(filename):
|
||
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
||
|
||
|
||
def get_timestamp():
|
||
return datetime.now().strftime('%y%m%d-%H%M%S')
|
||
|
||
|
||
def imshow(x, title=None, cbar=False, figsize=None):
|
||
plt.figure(figsize=figsize)
|
||
plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
|
||
if title:
|
||
plt.title(title)
|
||
if cbar:
|
||
plt.colorbar()
|
||
plt.show()
|
||
|
||
|
||
def surf(Z, cmap='rainbow', figsize=None):
|
||
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)
|
||
plt.show()
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# get image pathes
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
def get_image_paths(dataroot):
|
||
paths = None # return None if dataroot is None
|
||
if dataroot is not None:
|
||
paths = sorted(_get_paths_from_images(dataroot))
|
||
return paths
|
||
|
||
|
||
def _get_paths_from_images(path):
|
||
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)
|
||
return images
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# split large images into small images
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
def patches_from_image(img, p_size=512, p_overlap=64, p_max=800):
|
||
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)
|
||
|
||
return patches
|
||
|
||
|
||
def imssave(imgs, img_path):
|
||
"""
|
||
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.join(os.path.dirname(img_path), img_name+str('_s{:04d}'.format(i))+'.png')
|
||
cv2.imwrite(new_path, img)
|
||
|
||
|
||
def split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=800, p_overlap=96, p_max=1000):
|
||
"""
|
||
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, os.path.join(taget_dataroot,os.path.basename(img_path)))
|
||
#if original_dataroot == taget_dataroot:
|
||
#del img_path
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# makedir
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
def mkdir(path):
|
||
if not os.path.exists(path):
|
||
os.makedirs(path)
|
||
|
||
|
||
def mkdirs(paths):
|
||
if isinstance(paths, str):
|
||
mkdir(paths)
|
||
else:
|
||
for path in paths:
|
||
mkdir(path)
|
||
|
||
|
||
def mkdir_and_rename(path):
|
||
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)
|
||
os.makedirs(path)
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# read image from path
|
||
# opencv is fast, but read BGR numpy image
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
# --------------------------------------------
|
||
# get uint8 image of size HxWxn_channles (RGB)
|
||
# --------------------------------------------
|
||
def imread_uint(path, n_channels=3):
|
||
# 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
|
||
return img
|
||
|
||
|
||
# --------------------------------------------
|
||
# matlab's imwrite
|
||
# --------------------------------------------
|
||
def imsave(img, img_path):
|
||
img = np.squeeze(img)
|
||
if img.ndim == 3:
|
||
img = img[:, :, [2, 1, 0]]
|
||
cv2.imwrite(img_path, img)
|
||
|
||
def imwrite(img, img_path):
|
||
img = np.squeeze(img)
|
||
if img.ndim == 3:
|
||
img = img[:, :, [2, 1, 0]]
|
||
cv2.imwrite(img_path, img)
|
||
|
||
|
||
|
||
# --------------------------------------------
|
||
# get single image of size HxWxn_channles (BGR)
|
||
# --------------------------------------------
|
||
def read_img(path):
|
||
# 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]
|
||
return img
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# image format conversion
|
||
# --------------------------------------------
|
||
# numpy(single) <---> numpy(unit)
|
||
# numpy(single) <---> tensor
|
||
# numpy(unit) <---> tensor
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
# --------------------------------------------
|
||
# numpy(single) [0, 1] <---> numpy(unit)
|
||
# --------------------------------------------
|
||
|
||
|
||
def uint2single(img):
|
||
|
||
return np.float32(img/255.)
|
||
|
||
|
||
def single2uint(img):
|
||
|
||
return np.uint8((img.clip(0, 1)*255.).round())
|
||
|
||
|
||
def uint162single(img):
|
||
|
||
return np.float32(img/65535.)
|
||
|
||
|
||
def single2uint16(img):
|
||
|
||
return np.uint16((img.clip(0, 1)*65535.).round())
|
||
|
||
|
||
# --------------------------------------------
|
||
# numpy(unit) (HxWxC or HxW) <---> tensor
|
||
# --------------------------------------------
|
||
|
||
|
||
# convert uint to 4-dimensional torch tensor
|
||
def uint2tensor4(img):
|
||
if img.ndim == 2:
|
||
img = np.expand_dims(img, axis=2)
|
||
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0)
|
||
|
||
|
||
# convert uint to 3-dimensional torch tensor
|
||
def uint2tensor3(img):
|
||
if img.ndim == 2:
|
||
img = np.expand_dims(img, axis=2)
|
||
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.)
|
||
|
||
|
||
# convert 2/3/4-dimensional torch tensor to uint
|
||
def tensor2uint(img):
|
||
img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
|
||
if img.ndim == 3:
|
||
img = np.transpose(img, (1, 2, 0))
|
||
return np.uint8((img*255.0).round())
|
||
|
||
|
||
# --------------------------------------------
|
||
# numpy(single) (HxWxC) <---> tensor
|
||
# --------------------------------------------
|
||
|
||
|
||
# convert single (HxWxC) to 3-dimensional torch tensor
|
||
def single2tensor3(img):
|
||
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float()
|
||
|
||
|
||
# convert single (HxWxC) to 4-dimensional torch tensor
|
||
def single2tensor4(img):
|
||
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0)
|
||
|
||
|
||
# convert torch tensor to single
|
||
def tensor2single(img):
|
||
img = img.data.squeeze().float().cpu().numpy()
|
||
if img.ndim == 3:
|
||
img = np.transpose(img, (1, 2, 0))
|
||
|
||
return img
|
||
|
||
# convert torch tensor to single
|
||
def tensor2single3(img):
|
||
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)
|
||
return img
|
||
|
||
|
||
def single2tensor5(img):
|
||
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float().unsqueeze(0)
|
||
|
||
|
||
def single32tensor5(img):
|
||
return torch.from_numpy(np.ascontiguousarray(img)).float().unsqueeze(0).unsqueeze(0)
|
||
|
||
|
||
def single42tensor4(img):
|
||
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float()
|
||
|
||
|
||
# from skimage.io import imread, imsave
|
||
def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
|
||
'''
|
||
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.
|
||
return img_np.astype(out_type)
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# Augmentation, flipe and/or rotate
|
||
# --------------------------------------------
|
||
# The following two are enough.
|
||
# (1) augmet_img: numpy image of WxHxC or WxH
|
||
# (2) augment_img_tensor4: tensor image 1xCxWxH
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
def augment_img(img, mode=0):
|
||
'''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:
|
||
return np.flipud(np.rot90(img, k=3))
|
||
|
||
|
||
def augment_img_tensor4(img, mode=0):
|
||
'''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:
|
||
return img.rot90(3, [2, 3]).flip([2])
|
||
|
||
|
||
def augment_img_tensor(img, mode=0):
|
||
'''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)
|
||
|
||
return img_tensor.type_as(img)
|
||
|
||
|
||
def augment_img_np3(img, mode=0):
|
||
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)
|
||
return img
|
||
|
||
|
||
def augment_imgs(img_list, hflip=True, rot=True):
|
||
# 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
|
||
|
||
return [_augment(img) for img in img_list]
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# modcrop and shave
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
def modcrop(img_in, scale):
|
||
# 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))
|
||
return img
|
||
|
||
|
||
def shave(img_in, border=0):
|
||
# img_in: Numpy, HWC or HW
|
||
img = np.copy(img_in)
|
||
h, w = img.shape[:2]
|
||
img = img[border:h-border, border:w-border]
|
||
return img
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# image processing process on numpy image
|
||
# channel_convert(in_c, tar_type, img_list):
|
||
# rgb2ycbcr(img, only_y=True):
|
||
# bgr2ycbcr(img, only_y=True):
|
||
# ycbcr2rgb(img):
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
def rgb2ycbcr(img, only_y=True):
|
||
'''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.
|
||
return rlt.astype(in_img_type)
|
||
|
||
|
||
def ycbcr2rgb(img):
|
||
'''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.
|
||
return rlt.astype(in_img_type)
|
||
|
||
|
||
def bgr2ycbcr(img, only_y=True):
|
||
'''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.
|
||
return rlt.astype(in_img_type)
|
||
|
||
|
||
def channel_convert(in_c, tar_type, img_list):
|
||
# 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:
|
||
return img_list
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# metric, PSNR and SSIM
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
# --------------------------------------------
|
||
# PSNR
|
||
# --------------------------------------------
|
||
def calculate_psnr(img1, img2, border=0):
|
||
# 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')
|
||
return 20 * math.log10(255.0 / math.sqrt(mse))
|
||
|
||
|
||
# --------------------------------------------
|
||
# SSIM
|
||
# --------------------------------------------
|
||
def calculate_ssim(img1, img2, border=0):
|
||
'''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:
|
||
raise ValueError('Wrong input image dimensions.')
|
||
|
||
|
||
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()
|
||
|
||
|
||
'''
|
||
# --------------------------------------------
|
||
# matlab's bicubic imresize (numpy and torch) [0, 1]
|
||
# --------------------------------------------
|
||
'''
|
||
|
||
|
||
# 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))
|
||
|
||
|
||
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)
|
||
|
||
|
||
# --------------------------------------------
|
||
# imresize for tensor image [0, 1]
|
||
# --------------------------------------------
|
||
def imresize(img, scale, antialiasing=True):
|
||
# 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_()
|
||
return out_2
|
||
|
||
|
||
# --------------------------------------------
|
||
# imresize for numpy image [0, 1]
|
||
# --------------------------------------------
|
||
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()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
print('---')
|
||
# img = imread_uint('test.bmp', 3)
|
||
# img = uint2single(img)
|
||
# img_bicubic = imresize_np(img, 1/4) |