# 盲超分中Bicubic对不同模糊核sigma一次次保存结果再算有点麻烦，直接写到一起来测试
##1.定义模糊核下采样的过程，并用Bicubic上采样LR (low-resolution) image
- 下载pca_matrix.pth从DAN仓库greatlog/DAN: This is an official implementation of Unfolding the Alternating Optimization for Blind Super Resolution (github.com)
import os
import random
import cv2 as cv
import utils_image
import degradation
import torch
import data.img_file as imf
import numpy as np
class ImageDataset:
def __init__(self, lr_data_dir, hr_data_dir):
self.lr_data_dir = lr_data_dir
self.hr_data_dir = hr_data_dir
self.lr_paths = []
self.hr_paths = []
self._load_data()
def _load_data(self):
# 加载LR图像路径
for root, dirs, files in os.walk(self.lr_data_dir):
for file in files:
if file.endswith(".png"): #
lr_path = os.path.join(root, file)
self.lr_paths.append(lr_path)
# 加载HR图像路径
for root, dirs, files in os.walk(self.hr_data_dir):
for file in files:
if file.endswith(".png"): #
hr_path = os.path.join(root, file)
self.hr_paths.append(hr_path)
def get_random_pair(self):
index = random.randint(0, len(self.lr_paths) - 1)
lr_img = cv.imread(self.lr_paths[index])
hr_img = cv.imread(self.hr_paths[index])
return lr_img, hr_img
def get_all_pairs(self):
pairs = []
for lr_path, hr_path in zip(self.lr_paths, self.hr_paths):
lr_img = cv.imread(lr_path)
hr_img = cv.imread(hr_path)
pairs.append((lr_img, hr_img))
return pairs
def get_hr(self):
HR = []
for hr_path in self.hr_paths:
hr_img = cv.imread(hr_path)
HR.append(hr_img)
return HR
# load
urban = r"D:\Dataset\urban100\urban100"
set5 = r"D:\Dataset\Set5\Set5\original"
set14 = r"D:\Dataset\Set14\Set14\original"
bsd100 = r'D:\Dataset\BSD100\BSDS100'
m109 = r'D:\Dataset\manga109'
hr_data_dir = m109
lr_data_dir = r"D:\Dataset\Set5\Set5\LRbicx4"
dataset = ImageDataset(lr_data_dir, hr_data_dir)
HR_img = dataset.get_hr()
# pca_matrix
pca_matrix = torch.load(
os.path.abspath(r"C:\Users\yours\PycharmProjects\Query\pca_matrix.pth"), map_location=lambda storage, loc: storage
)
print("PCA matrix shape: {}".format(pca_matrix.shape))
up_scale = 6
mod_scale = 6
sig_3 = [0.2,0.4,0.6,0.8,1.0,1.35,1.5,1.65,1.8,1.95,2.10,2.25,2.4]
sig_4 = [1.8,2.0,2.2,2.4,2.6,2.8,3.0,3.2]
sig_6 = [3.8,4.0,4.2,4.4,4.6,4.8,5.0,5.2]
degradation_setting = {
"random_kernel": False,
"code_length": 10,
"ksize": 21,
"pca_matrix": pca_matrix,
"scale": up_scale,
"cuda": True,
"rate_iso": 1.0
}
PSNR_ALL = 0
SSIM_ALL = 0
for sig in sig_6:
print(sig)
count = 1
PSNR_all = 0
SSIM_all = 0
for hr_img in HR_img:
# C, H, W = hr_img.size()
width = int(np.floor(hr_img.shape[1] / mod_scale))
height = int(np.floor(hr_img.shape[0] / mod_scale))
# modcrop
if len(hr_img.shape) == 3:
hr_img = hr_img[0: mod_scale * height, 0: mod_scale * width, :]
else:
hr_img = hr_img[0: mod_scale * height, 0: mod_scale * width]
img_HR = imf.img2tensor(hr_img)
C, H, W = img_HR.size()
img_HR_old = imf.tensor2img(img_HR)
img_HR_old_Y = utils_image.bgr2ycbcr(img_HR_old)
prepro = degradation.SRMDPreprocessing(sig=sig,**degradation_setting)
LR_img, ker_map = prepro(img_HR.view(1, C, H, W))
image_LR_blur = imf.tensor2img(LR_img)
image_SR_bicubic = cv.resize(image_LR_blur, (W,H), interpolation=cv.INTER_CUBIC)
# cv.imwrite(os.path.join(r'C:\Users\yours\PycharmProjects\Query\result', 'sig{}_{}.png'.format(sig,count)), image_SR_bicubic)
image_SR_Y = utils_image.bgr2ycbcr(image_SR_bicubic)
# print("-----Processing on {}.image-----.".format(count))
# count = count + 1
PSNR = imf.calculate_psnr(img_HR_old,image_SR_bicubic)
SSIM = imf.calculate_ssim(img_HR_old,image_SR_bicubic)
# print(PSNR)
# psnr and ssim metrics
PSNR_all = PSNR + PSNR_all
SSIM_all = SSIM + SSIM_all
psnr_aver = PSNR_all / len(HR_img)
ssim_aver = SSIM_all / len(HR_img)
PSNR_ALL = PSNR_ALL + psnr_aver
SSIM_ALL = SSIM_ALL + ssim_aver
# print(len(sig_3))
PSNR_ALL_a = PSNR_ALL / len(sig_6)
SSIM_ALL_a = SSIM_ALL / len(sig_6)
print("Average PSNR and SSIM is {:2f} dB and {:2f} on scale factor {}".format(PSNR_ALL_a,SSIM_ALL_a, up_scale))
## 2、导入一些必要的文件从DAN（模糊核估计）类似的方法中
### 新建data文件夹，在下面如下的代码img_file.py
import math
import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision.utils import make_grid
try:
import accimage
except ImportError:
accimage = None
def _is_pil_image(img):
if accimage is not None:
return isinstance(img, (Image.Image, accimage.Image))
else:
return isinstance(img, Image.Image)
def _is_tensor_image(img):
return torch.is_tensor(img) and img.ndimension() == 3
def _is_numpy_image(img):
return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
def to_pil_image(pic, mode=None):
if not (_is_numpy_image(pic) or _is_tensor_image(pic)):
raise TypeError("pic should be Tensor or ndarray. Got {}.".format(type(pic)))
npimg = pic
if isinstance(pic, torch.FloatTensor):
pic = pic.mul(255).byte()
if torch.is_tensor(pic):
npimg = np.transpose(pic.numpy(), (1, 2, 0))
if not isinstance(npimg, np.ndarray):
raise TypeError(
"Input pic must be a torch.Tensor or NumPy ndarray, "
+ "not {}".format(type(npimg))
)
if npimg.shape[2] == 1:
expected_mode = None
npimg = npimg[:, :, 0]
if npimg.dtype == np.uint8:
expected_mode = "L"
if npimg.dtype == np.int16:
expected_mode = "I;16"
if npimg.dtype == np.int32:
expected_mode = "I"
elif npimg.dtype == np.float32:
expected_mode = "F"
if mode is not None and mode != expected_mode:
raise ValueError(
"Incorrect mode ({}) supplied for input type {}. Should be {}".format(
mode, np.dtype, expected_mode
)
)
mode = expected_mode
elif npimg.shape[2] == 4:
permitted_4_channel_modes = ["RGBA", "CMYK"]
if mode is not None and mode not in permitted_4_channel_modes:
raise ValueError(
"Only modes {} are supported for 4D inputs".format(
permitted_4_channel_modes
)
)
if mode is None and npimg.dtype == np.uint8:
mode = "RGBA"
else:
permitted_3_channel_modes = ["RGB", "YCbCr", "HSV"]
if mode is not None and mode not in permitted_3_channel_modes:
raise ValueError(
"Only modes {} are supported for 3D inputs".format(
permitted_3_channel_modes
)
)
if mode is None and npimg.dtype == np.uint8:
mode = "RGB"
if mode is None:
raise TypeError("Input type {} is not supported".format(npimg.dtype))
return Image.fromarray(npimg, mode=mode)
def to_tensor(pic):
if not (_is_pil_image(pic) or _is_numpy_image(pic)):
raise TypeError("pic should be PIL Image or ndarray. Got {}".format(type(pic)))
if isinstance(pic, np.ndarray):
# handle numpy array
img = torch.from_numpy(pic.transpose((2, 0, 1)))
# backward compatibility
return img.float().div(255)
if accimage is not None and isinstance(pic, accimage.Image):
nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)
pic.copyto(nppic)
return torch.from_numpy(nppic)
# handle PIL Image
if pic.mode == "I":
img = torch.from_numpy(np.array(pic, np.int32, copy=False))
elif pic.mode == "I;16":
img = torch.from_numpy(np.array(pic, np.int16, copy=False))
else:
img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
# PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
if pic.mode == "YCbCr":
nchannel = 3
elif pic.mode == "I;16":
nchannel = 1
else:
nchannel = len(pic.mode)
img = img.view(pic.size[1], pic.size[0], nchannel)
# put it from HWC to CHW format
# yikes, this transpose takes 80% of the loading time/CPU
img = img.transpose(0, 1).transpose(0, 2).contiguous()
if isinstance(img, torch.ByteTensor):
return img.float().div(255)
else:
return img
def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
"""
Converts a torch Tensor into an image Numpy array
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) # 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)
def save_img(img, img_path, mode="RGB"):
cv2.imwrite(img_path, img)
def img2tensor(img):
"""
# BGR to RGB, HWC to CHW, numpy to tensor
Input: img(H, W, C), [0,255], np.uint8 (default)
Output: 3D(C,H,W), RGB order, float tensor
"""
img = img.astype(np.float32) / 255.0
img = img[:, :, [2, 1, 0]]
img = torch.from_numpy(np.ascontiguousarray(np.transpose(img, (2, 0, 1)))).float()
return img
def calculate_psnr(img1, img2):
# img1 and img2 have range [0, 255]
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))
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(img1, img2):
"""calculate SSIM
the same outputs as MATLAB's
img1, img2: [0, 255]
"""
if not img1.shape == img2.shape:
raise ValueError("Input images must have the same dimensions.")
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, img2))
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.")
### util.py
import math
import os
import pickle
import random
import cv2
import numpy as np
import torch
# Files & IO
IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP']
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def _get_paths_from_images(path):
'''get image path list from image folder'''
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
def _get_paths_from_lmdb(dataroot):
'''get image path list from lmdb meta info'''
meta_info = pickle.load(open(os.path.join(dataroot, 'meta_info.pkl'), 'rb'))
paths = meta_info['keys']
sizes = meta_info['resolution']
if len(sizes) == 1:
sizes = sizes * len(paths)
return paths, sizes
def get_image_paths(data_type, dataroot):
'''get image path list
support lmdb or image files'''
paths, sizes = None, None
if dataroot is not None:
if data_type == 'lmdb':
paths, sizes = _get_paths_from_lmdb(dataroot)
return paths, sizes
elif data_type == 'img':
paths = sorted(_get_paths_from_images(dataroot))
return paths
else:
raise NotImplementedError('data_type [{:s}] is not recognized.'.format(data_type))
def _read_img_lmdb(env, key, size):
'''read image from lmdb with key (w/ and w/o fixed size)
size: (C, H, W) tuple'''
with env.begin(write=False) as txn:
buf = txn.get(key.encode('ascii'))
img_flat = np.frombuffer(buf, dtype=np.uint8)
C, H, W = size
img = img_flat.reshape(H, W, C)
return img
def read_img(env, path, size=None):
'''read image by cv2 or from lmdb
return: Numpy float32, HWC, BGR, [0,1]'''
if env is None: # img
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
else:
img = _read_img_lmdb(env, path, size)
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 processing
# process on numpy image
def augment(img, hflip=True, rot=True, mode=None):
# 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
if mode == 'LQ' or mode == 'GT' or mode == 'SRker':
return _augment(img)
elif mode == 'LQGTker':
return [_augment(I) for I in img]
def augment_flow(img_list, flow_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
def _augment_flow(flow):
if hflip:
flow = flow[:, ::-1, :]
flow[:, :, 0] *= -1
if vflip:
flow = flow[::-1, :, :]
flow[:, :, 1] *= -1
if rot90:
flow = flow.transpose(1, 0, 2)
flow = flow[:, :, [1, 0]]
return flow
rlt_img_list = [_augment(img) for img in img_list]
rlt_flow_list = [_augment_flow(flow) for flow in flow_list]
return rlt_img_list, rlt_flow_list
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
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 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 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 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
# Functions
# matlab 'imresize' function, now only support 'bicubic'
def cubic(x):
absx = torch.abs(x)
absx2 = absx ** 2
absx3 = absx ** 3
weight = (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))
return weight
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)
def imresize(img, scale, antialiasing=True):
# Now the scale should be the same for H and W
# input: img: CHW RGB [0,1]
# output: CHW RGB [0,1] w/o round
is_numpy = False
if isinstance(img, np.ndarray):
img = torch.from_numpy(img.transpose(2, 0, 1))
is_numpy = True
device = img.device
is_batch = True
if len(img.shape) == 3: # C, H, W
img = img[None]
is_batch = False
B, in_C, in_H, in_W = img.size()
img = img.view(-1, in_H, in_W)
_, 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_H, indices_H = weights_H.to(device), indices_H.to(device)
weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
in_W, out_W, scale, kernel, kernel_width, antialiasing)
weights_W, indices_W = weights_W.to(device), indices_W.to(device)
# process H dimension
# symmetric copying
img_aug = torch.FloatTensor(B * in_C, in_H + sym_len_Hs + sym_len_He, in_W).to(device)
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().to(device)
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().to(device)
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(B * in_C, out_H, in_W).to(device)
kernel_width = weights_H.size(1)
for i in range(out_H):
idx = int(indices_H[i][0])
out_1[:, i, :] = (img_aug[:, idx:idx + kernel_width, :].transpose(1, 2).matmul(
weights_H[i][None, :, None].repeat(B * in_C, 1, 1))).squeeze()
# process W dimension
# symmetric copying
out_1_aug = torch.FloatTensor(B * in_C, out_H, in_W + sym_len_Ws + sym_len_We).to(device)
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().to(device)
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().to(device)
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(B * in_C, out_H, out_W).to(device)
kernel_width = weights_W.size(1)
for i in range(out_W):
idx = int(indices_W[i][0])
out_2[:, :, i] = (out_1_aug[:, :, idx:idx + kernel_width].matmul(
weights_W[i][None, :, None].repeat(B * in_C, 1, 1))).squeeze()
out_2 = out_2.contiguous().view(B, in_C, out_H, out_W)
if not is_batch:
out_2 = out_2[0]
return out_2.cpu().numpy().transpose(1, 2, 0) if is_numpy else out_2
### Load data kernel map ###
def load_ker_map_list(path):
real_ker_map_list = []
batch_kermap = torch.load(path)
size_kermap = batch_kermap.size()
m = size_kermap[0]
for i in range(m):
real_ker_map_list.append(batch_kermap[i])
return real_ker_map_list
if __name__ == '__main__':
# test imresize function
# read images
img = cv2.imread('test.png')
img = img * 1.0 / 255
img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
# imresize
scale = 1 / 4
import time
total_time = 0
for i in range(10):
start_time = time.time()
rlt = imresize(img, scale, antialiasing=True)
use_time = time.time() - start_time
total_time += use_time
print('average time: {}'.format(total_time / 10))
import torchvision.utils
torchvision.utils.save_image((rlt * 255).round() / 255, 'rlt.png', nrow=1, padding=0,
normalize=False)
### degradation.py : SRMDProcessing的类
import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from data.util import imresize
from scipy.io import loadmat
from torch.autograd import Variable
def DUF_downsample(x, scale=4):
"""Downsamping with Gaussian kernel used in the DUF official code
Args:
x (Tensor, [B, T, C, H, W]): frames to be downsampled.
scale (int): downsampling factor: 2 | 3 | 4.
"""
assert scale in [2, 3, 4], "Scale [{}] is not supported".format(scale)
def gkern(kernlen=13, nsig=1.6):
import scipy.ndimage.filters as fi
inp = np.zeros((kernlen, kernlen))
# set element at the middle to one, a dirac delta
inp[kernlen // 2, kernlen // 2] = 1
# gaussian-smooth the dirac, resulting in a gaussian filter mask
return fi.gaussian_filter(inp, nsig)
B, T, C, H, W = x.size()
x = x.view(-1, 1, H, W)
pad_w, pad_h = 6 + scale * 2, 6 + scale * 2 # 6 is the pad of the gaussian filter
r_h, r_w = 0, 0
if scale == 3:
r_h = 3 - (H % 3)
r_w = 3 - (W % 3)
x = F.pad(x, [pad_w, pad_w + r_w, pad_h, pad_h + r_h], "reflect")
gaussian_filter = (
torch.from_numpy(gkern(13, 0.4 * scale)).type_as(x).unsqueeze(0).unsqueeze(0)
)
x = F.conv2d(x, gaussian_filter, stride=scale)
x = x[:, :, 2:-2, 2:-2]
x = x.view(B, T, C, x.size(2), x.size(3))
return x
def PCA(data, k=2):
X = torch.from_numpy(data)
X_mean = torch.mean(X, 0)
X = X - X_mean.expand_as(X)
U, S, V = torch.svd(torch.t(X))
return U[:, :k] # PCA matrix
def random_batch_kernel(
batch,
l=21,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
tensor=True,
random_disturb=False,
):
if rate_iso == 1:
sigma = np.random.uniform(sig_min, sig_max, (batch, 1, 1))
ax = np.arange(-l // 2 + 1.0, l // 2 + 1.0)
xx, yy = np.meshgrid(ax, ax)
xx = xx[None].repeat(batch, 0)
yy = yy[None].repeat(batch, 0)
kernel = np.exp(-(xx ** 2 + yy ** 2) / (2.0 * sigma ** 2))
kernel = kernel / np.sum(kernel, (1, 2), keepdims=True)
return torch.FloatTensor(kernel) if tensor else kernel
else:
sigma_x = np.random.uniform(sig_min, sig_max, (batch, 1, 1))
sigma_y = np.random.uniform(sig_min, sig_max, (batch, 1, 1))
D = np.zeros((batch, 2, 2))
D[:, 0, 0] = sigma_x.squeeze() ** 2
D[:, 1, 1] = sigma_y.squeeze() ** 2
radians = np.random.uniform(-np.pi, np.pi, (batch))
mask_iso = np.random.uniform(0, 1, (batch)) < rate_iso
radians[mask_iso] = 0
sigma_y[mask_iso] = sigma_x[mask_iso]
U = np.zeros((batch, 2, 2))
U[:, 0, 0] = np.cos(radians)
U[:, 0, 1] = -np.sin(radians)
U[:, 1, 0] = np.sin(radians)
U[:, 1, 1] = np.cos(radians)
sigma = np.matmul(U, np.matmul(D, U.transpose(0, 2, 1)))
ax = np.arange(-l // 2 + 1.0, l // 2 + 1.0)
xx, yy = np.meshgrid(ax, ax)
xy = np.hstack((xx.reshape((l * l, 1)), yy.reshape(l * l, 1))).reshape(l, l, 2)
xy = xy[None].repeat(batch, 0)
inverse_sigma = np.linalg.inv(sigma)[:, None, None]
kernel = np.exp(
-0.5
* np.matmul(
np.matmul(xy[:, :, :, None], inverse_sigma), xy[:, :, :, :, None]
)
)
kernel = kernel.reshape(batch, l, l)
if random_disturb:
kernel = kernel + np.random.uniform(0, 0.25, (batch, l, l)) * kernel
kernel = kernel / np.sum(kernel, (1, 2), keepdims=True)
return torch.FloatTensor(kernel) if tensor else kernel
def stable_batch_kernel(batch, l=21, sig=2.6, tensor=True):
sigma = sig
ax = np.arange(-l // 2 + 1.0, l // 2 + 1.0)
xx, yy = np.meshgrid(ax, ax)
xx = xx[None].repeat(batch, 0)
yy = yy[None].repeat(batch, 0)
kernel = np.exp(-(xx ** 2 + yy ** 2) / (2.0 * sigma ** 2))
kernel = kernel / np.sum(kernel, (1, 2), keepdims=True)
return torch.FloatTensor(kernel) if tensor else kernel
def b_Bicubic(variable, scale):
B, C, H, W = variable.size()
H_new = int(H / scale)
W_new = int(W / scale)
tensor_v = variable.view((B, C, H, W))
re_tensor = imresize(tensor_v, 1 / scale)
return re_tensor
def random_batch_noise(batch, high, rate_cln=1.0):
noise_level = np.random.uniform(size=(batch, 1)) * high
noise_mask = np.random.uniform(size=(batch, 1))
noise_mask[noise_mask < rate_cln] = 0
noise_mask[noise_mask >= rate_cln] = 1
return noise_level * noise_mask
def b_GaussianNoising(tensor, sigma, mean=0.0, noise_size=None, min=0.0, max=1.0):
if noise_size is None:
size = tensor.size()
else:
size = noise_size
noise = torch.mul(
torch.FloatTensor(np.random.normal(loc=mean, scale=1.0, size=size)),
sigma.view(sigma.size() + (1, 1)),
).to(tensor.device)
return torch.clamp(noise + tensor, min=min, max=max)
def b_GaussianNoising(tensor, noise_high, mean=0.0, noise_size=None, min=0.0, max=1.0):
if noise_size is None:
size = tensor.size()
else:
size = noise_size
noise = torch.FloatTensor(
np.random.normal(loc=mean, scale=noise_high, size=size)
).to(tensor.device)
return torch.clamp(noise + tensor, min=min, max=max)
class BatchSRKernel(object):
def __init__(
self,
l=21,
sig=2.6,
sig_min=0.2,
sig_max=4.0,
rate_iso=1.0,
random_disturb=False,
):
self.l = l
self.sig = sig
self.sig_min = sig_min
self.sig_max = sig_max
self.rate = rate_iso
self.random_disturb = random_disturb
def __call__(self, random, batch, tensor=False):
if random == True: # random kernel
return random_batch_kernel(
batch,
l=self.l,
sig_min=self.sig_min,
sig_max=self.sig_max,
rate_iso=self.rate,
tensor=tensor,
random_disturb=self.random_disturb,
)
else: # stable kernel
return stable_batch_kernel(batch, l=self.l, sig=self.sig, tensor=tensor)
class BatchBlurKernel(object):
def __init__(self, kernels_path):
kernels = loadmat(kernels_path)["kernels"]
self.num_kernels = kernels.shape[0]
self.kernels = kernels
def __call__(self, random, batch, tensor=False):
index = np.random.randint(0, self.num_kernels, batch)
kernels = self.kernels[index]
return torch.FloatTensor(kernels).contiguous() if tensor else kernels
class PCAEncoder(nn.Module):
def __init__(self, weight):
super().__init__()
self.register_buffer("weight", weight)
self.size = self.weight.size()
def forward(self, batch_kernel):
B, H, W = batch_kernel.size() # [B, l, l]
return torch.bmm(
batch_kernel.view((B, 1, H * W)), self.weight.expand((B,) + self.size)
).view((B, -1))
class BatchBlur(object):
def __init__(self, l=15):
self.l = l
if l % 2 == 1:
self.pad = (l // 2, l // 2, l // 2, l // 2)
else:
self.pad = (l // 2, l // 2 - 1, l // 2, l // 2 - 1)
# self.pad = nn.ZeroPad2d(l // 2)
def __call__(self, input, kernel):
B, C, H, W = input.size()
pad = F.pad(input, self.pad, mode='reflect')
H_p, W_p = pad.size()[-2:]
if len(kernel.size()) == 2:
input_CBHW = pad.view((C * B, 1, H_p, W_p))
kernel_var = kernel.contiguous().view((1, 1, self.l, self.l))
return F.conv2d(input_CBHW, kernel_var, padding=0).view((B, C, H, W))
else:
input_CBHW = pad.view((1, C * B, H_p, W_p))
kernel_var = (
kernel.contiguous()
.view((B, 1, self.l, self.l))
.repeat(1, C, 1, 1)
.view((B * C, 1, self.l, self.l))
)
return F.conv2d(input_CBHW, kernel_var, groups=B * C).view((B, C, H, W))
class SRMDPreprocessing(object):
def __init__(
self, scale, pca_matrix,
ksize=21, code_length=10,
random_kernel=True, noise=False, cuda=False, random_disturb=False,
sig=0, sig_min=0, sig_max=0, rate_iso=1.0, rate_cln=1, noise_high=0,
stored_kernel=False, pre_kernel_path=None
):
self.encoder = PCAEncoder(pca_matrix).cuda() if cuda else PCAEncoder(pca)
self.kernel_gen = (
BatchSRKernel(
l=ksize,
sig=sig, sig_min=sig_min, sig_max=sig_max,
rate_iso=rate_iso, random_disturb=random_disturb,
) if not stored_kernel else
BatchBlurKernel(pre_kernel_path)
)
self.blur = BatchBlur(l=ksize)
self.para_in = code_length
self.l = ksize
self.noise = noise
self.scale = scale
self.cuda = cuda
self.rate_cln = rate_cln
self.noise_high = noise_high
self.random = random_kernel
def __call__(self, hr_tensor, kernel=False):
# hr_tensor is tensor, not cuda tensor
hr_var = Variable(hr_tensor).cuda() if self.cuda else Variable(hr_tensor)
device = hr_var.device
B, C, H, W = hr_var.size()
b_kernels = Variable(self.kernel_gen(self.random, B, tensor=True)).to(device)
hr_blured_var = self.blur(hr_var, b_kernels)
# B x self.para_input
kernel_code = self.encoder(b_kernels)
# Down sample
if self.scale != 1:
lr_blured_t = b_Bicubic(hr_blured_var, self.scale)
else:
lr_blured_t = hr_blured_var
# Noisy
if self.noise:
Noise_level = torch.FloatTensor(
random_batch_noise(B, self.noise_high, self.rate_cln)
)
lr_noised_t = b_GaussianNoising(lr_blured_t, self.noise_high)
else:
Noise_level = torch.zeros((B, 1))
lr_noised_t = lr_blured_t
Noise_level = Variable(Noise_level).cuda()
re_code = (
torch.cat([kernel_code, Noise_level * 10], dim=1)
if self.noise
else kernel_code
)
lr_re = Variable(lr_noised_t).to(device)
return (lr_re, re_code, b_kernels) if kernel else (lr_re, re_code)
扫一扫
3736
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
