def __getitem__(self, idx):
# load image
img_np = self.load_img(self.image_filenames[idx])
# original HR image size
h, w, _ = img_np.shape
# determine valid HR image size with scale factor
HR_img_w = self.calculate_valid_crop_size(w, self.scale_factor)
HR_img_h = self.calculate_valid_crop_size(h, self.scale_factor)
# determine lr_img LR image size
LR_img_w = HR_img_w // self.scale_factor
LR_img_h = HR_img_h // self.scale_factor
HR_img_np = img_np[:HR_img_h, :HR_img_w, :] # high resolution image
LR_img_np = imresize(HR_img_np, scalar_scale=1.0 /
self.scale_factor) # low resolution image
BC_img_np = imresize(
LR_img_np,
scalar_scale=self.scale_factor) # bicubic upsampled image
if self.upsample_LR_patch:
LR_img_np = BC_img_np
HR_img = torch.from_numpy(HR_img_np).type(torch.FloatTensor).permute(
2, 0, 1) # size : 3(c) x h x w
LR_img = torch.from_numpy(LR_img_np).type(torch.FloatTensor).permute(
2, 0, 1) # size : 3(c) x (h/scale_factor) x (w/scale_factor)
BC_img = torch.from_numpy(BC_img_np).type(torch.FloatTensor).permute(
2, 0, 1) # size : 3(c) x h x w
return HR_img, LR_img, BC_img # YCbCr images
def __getitem__(self, idx):
img_name, img_lr, img_hr = self.image_read(idx)
img_lr, img_hr = self.transforms(img_lr, img_hr)
if cfg.CONST.SCALE == 4:
img_lr_s = imresize(img_lr/255.0, 1.0/2)*255
else:
img_lr_s = imresize(img_lr/255.0, 1.0/cfg.CONST.SCALE)*255
img_lr_s = img_lr_s.clamp(0,255)
return img_name, img_lr_s, img_lr, img_hr
def loss_fn(pred_imgs, pred_scale, target_imgs, target_scales):
batch_size = len(target_scales)
# pred_imgs : [repeat, batch_size, 1, w, h]
# target_scales : [batch_size]
if type(pred_imgs) == list:
reconstruction_loss = 0
length = len(pred_imgs)
ratio = np.power(scale_factor, 1 / length)
scale_v = [np.power(ratio, length - i - 1) for i in range(length)]
t_s_v = [[] for i in range(batch_size)]
for i, ts in enumerate(target_scales):
for j in range(length):
tmp_ts = ts / np.power(ratio, j + 1)
t_s_v[i].append(tmp_ts)
if tmp_ts < 1:
break
for i, pred in enumerate(pred_imgs):
for j, p in enumerate(pred):
if i >= len(t_s_v[j]):
break
scale = 1 / t_s_v[j][i]
target_img = target_imgs[j].unsqueeze(0)
p = p.unsqueeze(0)
#t = UpsamplingBilinear2d(scale_factor=scale)(target_img)
#t = UpsamplingBilinear2d(size=(
# target_img.shape[-2], target_img.shape[-1]))(t)
target_img = target_img.squeeze(0)
#if j == 0:
# print(1/target_scales[j] , scale)
t = imresize(target_img.cpu().numpy(), scalar_scale=scale)
t = imresize(t,
output_shape=(target_img.shape[-2],
target_img.shape[-1]))
t = torch.Tensor(t).cuda().unsqueeze(0).unsqueeze(0)
reconstruction_loss += \
np.sqrt(length - i) * torch.nn.MSELoss(reduction=reduction)(p, t)
else:
reconstruction_loss = \
torch.nn.MSELoss(reduction=reduction)(pred_imgs, target_imgs)
target_scales = target_scales.view(batch_size, -1)
#scale_loss = torch.nn.MSELoss(reduction='mean')(pred_scale, target_scales)
one_hot_target_scales = nn.functional.one_hot(
(target_scales * 5 - 5).long(), 16).float()
scale_loss = torch.nn.BCELoss(reduction='mean')(
pred_scale, one_hot_target_scales.cuda())
total_loss = reconstruction_loss + 10 * scale_loss
return total_loss
def __getitem__(self, idx):
img_name, img_lr = self.image_read(idx)
img_lr = img_lr[:, :, [2, 1, 0]]
img_lr = img2tensor(img_lr)
if cfg.CONST.SCALE == 4:
img_lr_s = imresize(img_lr / 255.0, 1.0 / 2) * 255
else:
img_lr_s = imresize(img_lr / 255.0, 1.0 / cfg.CONST.SCALE) * 255
img_lr_s = img_lr_s.clamp(0, 255)
return img_name, img_lr_s, img_lr
def get_test_data(self):
data_path = self.config['data_path'] + self.config['test_path']
files_list = [
join(data_path, x) for x in sorted(listdir(data_path))
if is_image_file(x)
]
img_list = []
# TODO: use scipy.misc when tuning, finally use imresize for matlab-like bicubic performance
for file in files_list:
img = io.imread(file)
if img.shape[2] == 4:
img = color.rgba2rgb(img)
img_ycbcr = color.rgb2ycbcr(img) / 255
img_ycbcr = img_ycbcr.astype('float32')
(rows, cols, channel) = img_ycbcr.shape
img_y, img_cb, img_cr = np.split(img_ycbcr,
indices_or_sections=channel,
axis=2)
size_lr = (int(rows // self.config['upscale']),
int(cols // self.config['upscale']))
# img_y_lr = cv2.resize(img_y.squeeze(), dsize=(int(cols // self.config['upscale']), int(rows // self.config['upscale'])),
# interpolation=cv2.INTER_CUBIC)
# img_y_lr = misc.imresize(img_y.squeeze(), size=size_lr, interp='bicubic', mode='F')
img_y_lr = imresize(img_y.squeeze(), output_shape=size_lr)
img_cb = imresize(img_cb.squeeze(), output_shape=size_lr)
img_cr = imresize(img_cr.squeeze(), output_shape=size_lr)
img = imresize(img, output_shape=size_lr)
# import matplotlib.pyplot as plt
# plt.imshow(img_y_lr, cmap ='gray')
# plt.show()
img_bundle = {
'name': os.path.basename(file),
'origin': img,
'x': img_y_lr,
'y': img_ycbcr,
'cb': img_cb.squeeze(),
'cr': img_cr.squeeze(),
'size': img_ycbcr.shape[0:-1]
}
# img_bundle = {'name': os.path.basename(file), 'x': img_y_lr, 'y': img_y, 'cb': img_cb.squeeze(),
# 'cr': img_cr.squeeze(),
# 'size': size_lr}
img_list.append(img_bundle)
return img_list
def creat_reals_pyramid(real, reals, opt):
real = real[:, 0:3, :, :]
for i in range(0, opt.stop_scale + 1, 1):
scale = math.pow(opt.scale_factor, opt.stop_scale - i)
curr_real = imresize(real, scale, opt)
reals.append(curr_real)
return reals
def adjust_scales2image(real_, opt):
opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1)
real = imresize(real_, opt.scale1, opt)
# opt.stop_scale = opt.train_stages - 1
opt.scale_factor = math.pow(
opt.min_size / (min(real.shape[2], real.shape[3])), 1 / opt.stop_scale)
return real
def TuiGAN_transfer(Gs,
Zs,
reals,
NoiseAmp,
Gs2,
opt,
in_s=None,
gen_start_scale=0):
if in_s is None:
in_s = torch.full(reals[0].shape, 0, device=opt.device)
x_ab = in_s
x_aba = in_s
count = 0
dir2save = functions.generate_dir2save(opt)
try:
os.makedirs(dir2save)
except OSError:
pass
for G, G2, Z_opt, real_curr, real_next, noise_amp in zip(
Gs, Gs2, Zs, reals, reals[1:], NoiseAmp):
z = functions.generate_noise([3, Z_opt.shape[2], Z_opt.shape[3]],
device=opt.device)
z = z.expand(real_curr.shape[0], 3, z.shape[2], z.shape[3])
x_ab = x_ab[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]]
z_in = noise_amp * z + real_curr
x_ab = G(z_in.detach(), x_ab)
x_aba = G2(x_ab, x_aba)
x_ab = imresize(x_ab.detach(), 1 / opt.scale_factor, opt)
x_ab = x_ab[:, :, 0:real_next.shape[2], 0:real_next.shape[3]]
x_aba = imresize(x_aba.detach(), 1 / opt.scale_factor, opt)
x_aba = x_aba[:, :, 0:real_next.shape[2], 0:real_next.shape[3]]
count += 1
plt.imsave('%s/x_ab_%d.png' % (dir2save, count),
functions.convert_image_np(x_ab.detach()),
vmin=0,
vmax=1)
plt.imsave('%s.png' % (dir2save),
functions.convert_image_np(x_ab.detach()),
vmin=0,
vmax=1)
# plt.imsave('%s.jpg' % (dir2save), functions.convert_image_np(x_ab.detach()), vmin=0,vmax=1)
return x_ab.detach()
def create_reals_pyramid(real, opt):
reals = []
# use old rescaling method for harmonization
if opt.train_mode == "harmonization":
for i in range(opt.stop_scale):
scale = math.pow(opt.scale_factor, opt.stop_scale - i)
curr_real = imresize(real, scale, opt)
reals.append(curr_real)
# use new rescaling method for all other tasks
else:
for i in range(opt.stop_scale):
scale = math.pow(
opt.scale_factor,
((opt.stop_scale - 1) / math.log(opt.stop_scale)) *
math.log(opt.stop_scale - i) + 1)
curr_real = imresize(real, scale, opt)
reals.append(curr_real)
reals.append(real)
return reals
def two_masked_mixed_images(image1, image2):
"""
:param np.array image1:
:param np.array image2:
:return:
"""
a = np.random.rand(1, image1.shape[1] // 32, image1.shape[2] // 32) / 2
b = np.random.rand(1, image1.shape[1] // 32,
image1.shape[2] // 32) / 2 + 0.5
a = np.clip(
imresize(a.transpose(1, 2, 0),
output_shape=(image1.shape[1], image1.shape[2]),
kernel='linear').transpose(2, 0, 1), 0, 1)
b = np.clip(
imresize(b.transpose(1, 2, 0),
output_shape=(image1.shape[1], image1.shape[2]),
kernel='linear').transpose(2, 0, 1), 0, 1)
return a * image1 + (1 - a) * image2, b * image1 + (1 - b) * image2
def __getitem__(self, idx):
img_name, img_lr, img_hr = self.image_read(idx)
img_lr, img_hr = self.transforms(img_lr, img_hr)
_, h, w = img_lr.size()
img_lr_ = img_lr[:, :int(h -
h % self.down_scale), :int(w - w %
self.down_scale)]
img_lr_s = imresize(img_lr_ / 255.0, 1.0 / self.down_scale) * 255
img_lr_s = img_lr_s.clamp(0, 255)
return img_name, img_lr_s, img_lr, img_hr
def cycle_rec(Gs,Gs2,Zs,reals,NoiseAmp,in_s,m_noise,m_image,opt, epoch):
x_ab = in_s
x_aba = in_s
if len(Gs) > 0:
count = 0
for G,G2,Z_opt,real_curr,real_next,noise_amp in zip(Gs,Gs2,Zs,reals,reals[1:],NoiseAmp):
z = functions.generate_noise([3, Z_opt.shape[2] , Z_opt.shape[3] ], device=opt.device)
z = z.expand(opt.bsz, 3, z.shape[2], z.shape[3])
z = m_noise(z)
x_ab = x_ab[:,:,0:real_curr.shape[2],0:real_curr.shape[3]]
x_ab = m_image(x_ab)
z_in = noise_amp*z+m_image(real_curr)
x_ab = G(z_in.detach(),x_ab)
x_aba = G2(x_ab,x_aba)
x_ab = imresize(x_ab.detach(),1/opt.scale_factor,opt)
x_ab = x_ab[:,:,0:real_next.shape[2],0:real_next.shape[3]]
x_aba = imresize(x_aba.detach(),1/opt.scale_factor,opt)
x_aba = x_aba[:,:,0:real_next.shape[2],0:real_next.shape[3]]
count += 1
return x_ab, x_aba
def __getitem__(self, idx):
img_name, img_lr = self.image_read(idx)
img_lr = img_lr[:, :, [2, 1, 0]]
img_lr = img2tensor(img_lr)
_, h, w = img_lr.size()
img_lr_ = img_lr[:, :int(h -
h % self.down_scale), :int(w - w %
self.down_scale)]
img_lr_s = imresize(img_lr_ / 255.0, 1.0 / self.down_scale) * 255
img_lr_s = img_lr_s.clamp(0, 255)
return img_name, img_lr_s, img_lr
def draw_concat(Gs,Zs,reals,NoiseAmp,in_s,mode,m_noise,m_image,opt):
G_z = in_s
if len(Gs) > 0:
if mode == 'rec':
count = 0
for G,Z_opt,real_curr,real_next,noise_amp in zip(Gs,Zs,reals,reals[1:],NoiseAmp):
G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]]
G_z = m_image(G_z)
z_in = m_image(real_curr)
G_z = G(z_in.detach(),G_z)
G_z = imresize(G_z.detach(),1/opt.scale_factor,opt)
G_z = G_z[:,:,0:real_next.shape[2],0:real_next.shape[3]]
count += 1
return G_z
def __getitem__(self, idx):
HR_patch_np = self.HR_patches_np[idx] # high resolution patch
LR_patch_np = imresize(HR_patch_np,
scalar_scale=1.0 /
self.scale_factor) # low resolution patch
BC_patch_np = imresize(
LR_patch_np,
output_shape=HR_patch_np.shape[-2:]) # bicubic upsampled patch
if self.transform:
LR_patch_np = self.transform(LR_patch_np)
if self.upsample_LR_patch:
LR_patch_np = BC_patch_np
HR_patch = torch.from_numpy(HR_patch_np).type(torch.FloatTensor)
LR_patch = torch.from_numpy(LR_patch_np).type(torch.FloatTensor)
BC_patch = torch.from_numpy(BC_patch_np).type(torch.FloatTensor)
HR_patch = HR_patch.unsqueeze(0) # size : 1(c) x 64(h) x 64(w)
LR_patch = LR_patch.unsqueeze(0) # size : 1(c) x 16(h) x 16(w)
BC_patch = BC_patch.unsqueeze(0) # size : 1(c) x 64(h) x 64(w)
return HR_patch, LR_patch, BC_patch # Y-channel patches
def __getitem__(self, index):
# load image
# img = io.imread(self.image_filenames[index])
img = Image.open(self.image_filenames[index]).convert('RGB')
# original HR image size
# (hr_img_h, hr_img_w, channel) = img.shape
hr_img_w = img.size[0]
hr_img_h = img.size[1]
if self.is_gray:
# only Y-channel is super-resolved
# skimage can convert RGB to YCbCr in [16, 235] while PIL cannot.
img = np.asarray(img)
img = color.rgb2ycbcr(img) / 255
channel = len(img.shape)
img, _, _ = np.split(img, indices_or_sections=channel, axis=-1)
# img = img.convert('YCbCr')
# precision degrade from float64 to float32
img = Image.fromarray(img.squeeze())
# img, _, _ = img.split()
# img = Image.fromarray(img.squeeze())
# hr_img HR image
hr_transform = tfs.ToTensor()
hr_img = hr_transform(img)
# determine lr_img LR image size
lr_img_list = []
for sf in self.scale_factor:
lr_img_w = hr_img_w // sf
lr_img_h = hr_img_h // sf
# lr_img LR image
# tfs resize set size as (h, w)
img = np.asarray(img)
img = imresize(img, output_shape=(lr_img_h, lr_img_w))
img = Image.fromarray(img.squeeze())
lr_transform = tfs.ToTensor()
lr_img = lr_transform(img)
lr_img_list.append(lr_img)
# Bicubic interpolated image
# bc_transform = tfs.Compose([tfs.ToPILImage(), tfs.Resize((hr_img_w, hr_img_h), interpolation=Image.BICUBIC), tfs.ToTensor()])
# bc_img = bc_transform(lr_img)
return lr_img_list, hr_img
def get_imresize_downsampled(image, downsampling_factor, downsampling_number):
"""
image is of type np.array
downsampling_factor should be integer - e.g. 0.5
"""
# TODO: move kernel type to args
downsampled_images = [image]
for i in range(1, downsampling_number + 1):
im = np.clip(
imresize(
image.transpose(1, 2, 0),
scale_factor=(
1 -
(downsampling_factor * downsampling_number))).transpose(
2, 0, 1), 0, 1)
downsampled_images.append(pil_to_np(crop_image(np_to_pil(im), d=32)))
return downsampled_images
def preprocess(directory,
save_npy_path,
patch_size=64,
stride=64,
scale=[1, 0.9, 0.8, 0.7, 0.6, 2],
rotation=[0, 1, 2, 3],
flip=[True, False],
is_gray=True):
HR_SET = []
images = sorted(glob.glob(directory + "*.png"))
print("The number of training images : ", len(images))
for idx in range(len(images)):
print("\r Processing ", idx + 1, " / ", len(images), end='')
image_directory = images[idx]
for f in flip:
for r in rotation:
for s in scale:
# load image
image = load_img(image_directory,
is_gray) # is_gary : YCbCr or RGB
# flipping
if f:
image = np.fliplr(image)
# rotation
image = np.rot90(image, k=r, axes=(0, 1))
# scaling
image = imresize(image, scalar_scale=s)
image = image.clip(0, 1)
# generate HR patch
h, w, _ = image.shape
for i in range(0, h - patch_size, stride):
for j in range(0, w - patch_size, stride):
hr_patch = image[i:i + patch_size,
j:j + patch_size, :]
if is_gray:
hr_patch = hr_patch[:, :, 0]
HR_SET.append(hr_patch)
print("\nThe number of training patches : ", len(HR_SET))
np.save(save_npy_path, HR_SET)
print("Training patches are successfully saved")
def resize_tensor(tensor, size=None, scale=None):
# TODO: Use multiple interpolation methods to reduce bias to bicubic
# TODO: tackle 3-channels
# TODO: Change all the bicubic operators as matlab bicubic
if scale is not None:
hr_w = tensor.shape[2]
hr_h = tensor.shape[3]
lr_w = int(np.ceil(hr_w * scale))
lr_h = int(np.ceil(hr_h * scale))
size = (lr_w, lr_h)
# print(tensor.shape)
y_numpy = tensor.data.cpu().numpy().squeeze(axis=1)
# print(y_numpy.shape)
newt = [imresize.imresize(y_sub, output_shape=size) for y_sub in y_numpy]
# tensor = [misc.imresize(y_sub, size=size, interp='bicubic', mode='F') for y_sub in y_numpy]
newt = np.array(newt, dtype=np.float32)
newt = np.expand_dims(newt, axis=1)
newt = torch.from_numpy(newt)
return newt
def adjust_scales2image(real_, opt):
opt.num_scales = math.ceil((math.log(
math.pow(opt.min_size / (min(real_.shape[2], real_.shape[3])), 1),
opt.scale_factor_init))) + 1
scale2stop = math.ceil(
math.log(
min([opt.max_size,
max([real_.shape[2], real_.shape[3]])]) /
max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init))
opt.stop_scale = opt.num_scales - scale2stop
opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1)
real = imresize(real_, opt.scale1, opt)
opt.scale_factor = math.pow(
opt.min_size / (min(real.shape[2], real.shape[3])),
1 / (opt.stop_scale))
scale2stop = math.ceil(
math.log(
min([opt.max_size,
max([real_.shape[2], real_.shape[3]])]) /
max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init))
opt.stop_scale = opt.num_scales - scale2stop
return real
def __getitem__(self, index):
# load image
# img = io.imread(self.image_filenames[index])
if self.preload:
img = self.image_list[index]
else:
img = Image.open(self.image_filenames[index]).convert('RGB')
# original HR image size
# (hr_img_h, hr_img_w, channel) = img.shape
hr_img_w = img.size[0]
hr_img_h = img.size[1]
# random scaling between [0.5, 1.0]
if self.random_scale:
eps = 1e-3
# ratio = random.uniform(0.5, 1)
ratio = random.randint(5, 10) * 0.1
if hr_img_w * ratio < self.crop_size:
ratio = self.crop_size / hr_img_w + eps
if hr_img_h * ratio < self.crop_size:
ratio = self.crop_size / hr_img_h + eps
scale_w = int(hr_img_w * ratio)
scale_h = int(hr_img_h * ratio)
# tfs resize set size as (h, w)
# transform = tfs.Resize((scale_h, scale_w), interpolation=Image.BICUBIC)
img = np.asarray(img)
img = imresize(img, output_shape=(scale_h, scale_w))
img = Image.fromarray(img.squeeze())
# img = transform(img)
# random crop
if self.crop_size:
transform = tfs.RandomCrop(self.crop_size)
img = transform(img)
# random rotation between [90, 180, 270] degrees
if self.rotate:
rv = random.randint(0, 3)
img = img.rotate(90 * rv, expand=True)
# random horizontal flip
if self.fliplr:
transform = tfs.RandomHorizontalFlip()
img = transform(img)
# random vertical flip
if self.fliptb:
if random.random() < 0.5:
img = img.transpose(Image.FLIP_TOP_BOTTOM)
# REMEMBER: change to gray image in the last, otherwise the color range will be out of [16, 235]
if self.is_gray:
# only Y-channel is super-resolved
# skimage can convert RGB to YCbCr in [16, 235] while PIL cannot.
img = np.asarray(img)
img = color.rgb2ycbcr(img) / 255
channel = len(img.shape)
img, _, _ = np.split(img, indices_or_sections=channel, axis=-1)
# img = img.convert('YCbCr')
# precision degrade from float64 to float32
img = Image.fromarray(img.squeeze())
# img, _, _ = img.split()
# hr_img HR image
hr_transform = tfs.ToTensor()
img = hr_transform(img)
# easy = img.data.cpu().numpy()
# # easy = np.asarray(img)
# if easy.min() < 0.0627 or easy.max() > 0.922:
# print('ERROR IN LOADING Y CHANNEL!')
# Bicubic interpolated image
# bc_transform = tfs.Compose([tfs.ToPILImage(), tfs.Resize((hr_img_w, hr_img_h), interpolation=Image.BICUBIC), ToTensor()])
# bc_img = bc_transform(lr_img)
return img
def wmcnn_hdf5(path, full=False):
def expanddims(array, expand=2):
for i in range(expand):
array = np.expand_dims(array, 0)
return array
# from matplotlib import pyplot as plt
# from skimage import measure
scale = 2
if not full:
size_label = 96
stride = 48
else:
size_label = 400
size_input = size_label / scale
batch_size = 400
classes = 7
# downsizes = [1, 0.7, 0.5]
id_files = join(path, 'id_files.txt')
# id_files = join(path, 'test_id.txt')
with open(id_files, 'r') as f:
article = f.readlines()
fnum = len(article) // classes
order = np.random.permutation(fnum)
training = order[:int(fnum * 0.7)]
testing = order[len(training):]
img_list = []
for i, line in enumerate(article):
img = io.imread(path + line[2:].strip('\n'))
if i % fnum in testing:
cls = line.split('/')[1]
mkdir_if_not_exist(path + 'test/' + cls)
io.imsave(path + 'test/' + line[2:].strip('\n'), img)
else:
if img.shape[2] == 4:
img = color.rgba2rgb(img)
img_ycbcr = color.rgb2ycbcr(img) / 255
(rows, cols, channel) = img_ycbcr.shape
img_y, img_cb, img_cr = np.split(img_ycbcr,
indices_or_sections=channel,
axis=2)
img_list.append(img_y)
for idx, img_gt in enumerate(img_list):
if idx == 0:
hf = h5py.File('D:/data/rsscn7/full_wdata.h5', 'w')
d_data = hf.create_dataset("data",
(batch_size, 1, size_input, size_input),
maxshape=(None, 1, size_input,
size_input),
dtype='float32')
d_ca = hf.create_dataset("CA",
(batch_size, 1, size_input, size_input),
maxshape=(None, 1, size_input,
size_input),
dtype='float32')
d_ch = hf.create_dataset("CH",
(batch_size, 1, size_input, size_input),
maxshape=(None, 1, size_input,
size_input),
dtype='float32')
d_cv = hf.create_dataset("CV",
(batch_size, 1, size_input, size_input),
maxshape=(None, 1, size_input,
size_input),
dtype='float32')
d_cd = hf.create_dataset("CD",
(batch_size, 1, size_input, size_input),
maxshape=(None, 1, size_input,
size_input),
dtype='float32')
else:
hf = h5py.File('D:/data/rsscn7/full_wdata.h5', 'a')
d_data = hf['data']
d_ca = hf['CA']
d_ch = hf['CH']
d_cv = hf['CV']
d_cd = hf['CD']
count = 0
if not full:
d_data.resize([idx * batch_size + 392, 1, size_input, size_input])
d_ca.resize([idx * batch_size + 392, 1, size_input, size_input])
d_ch.resize([idx * batch_size + 392, 1, size_input, size_input])
d_cv.resize([idx * batch_size + 392, 1, size_input, size_input])
d_cd.resize([idx * batch_size + 392, 1, size_input, size_input])
for flip in range(2):
for degree in range(4):
# for downsize in downsizes:
img = img_gt.squeeze()
if flip == 1:
img = np.fliplr(img)
for turn in range(degree):
img = np.rot90(img)
# img = imresize(img, scalar_scale=downsize)
hei, wid = img.shape
# fig = plt.figure(figsize=(6, 3))
for x in range(0, hei - size_label, stride):
for y in range(0, wid - size_label, stride):
subim_label = img[x:x + size_label,
y:y + size_label]
subim_data = imresize(subim_label,
scalar_scale=1 / scale)
coeffs2 = pywt.dwt2(subim_label, 'bior1.1')
LL, (LH, HL, HH) = coeffs2
d_data[idx * batch_size + count] = expanddims(
subim_data, expand=2)
d_ca[idx * batch_size + count] = expanddims(
LL, expand=2)
d_ch[idx * batch_size + count] = expanddims(
LH, expand=2)
d_cv[idx * batch_size + count] = expanddims(
HL, expand=2)
d_cd[idx * batch_size + count] = expanddims(
HH, expand=2)
count += 1
else:
d_data.resize([idx * batch_size + 1, 1, size_input, size_input])
d_ca.resize([idx * batch_size + 1, 1, size_input, size_input])
d_ch.resize([idx * batch_size + 1, 1, size_input, size_input])
d_cv.resize([idx * batch_size + 1, 1, size_input, size_input])
d_cd.resize([idx * batch_size + 1, 1, size_input, size_input])
img = img_gt.squeeze()
im_data = imresize(img, scalar_scale=1 / scale)
coeffs2 = pywt.dwt2(img, 'bior1.1')
LL, (LH, HL, HH) = coeffs2
d_data[idx * batch_size + count] = expanddims(im_data, expand=2)
d_ca[idx * batch_size + count] = expanddims(LL, expand=2)
d_ch[idx * batch_size + count] = expanddims(LH, expand=2)
d_cv[idx * batch_size + count] = expanddims(HL, expand=2)
d_cd[idx * batch_size + count] = expanddims(HH, expand=2)
count += 1
batch_size = count
hf.close()
def train(opt,Gs,Zs,reals,NoiseAmp, Gs2,Zs2,reals2,NoiseAmp2):
real_, real_2 = functions.read_two_domains(opt)
in_s = 0
in_s2 = 0
scale_num = 0
real = imresize(real_,opt.scale1,opt)
real2 = imresize(real_2,opt.scale1,opt)
reals = functions.creat_reals_pyramid(real,reals,opt)
reals2 = functions.creat_reals_pyramid(real2,reals2,opt)
nfc_prev = 0
errD_plot = []
errD2_plot = []
errG_plot = []
errG2_plot = []
rec_loss_plot = []
rec_loss2_plot = []
cyc_loss_plot = []
cyc_loss2_plot = []
while scale_num<opt.stop_scale+1:
opt.nfc = min(opt.nfc_init * pow(2, math.floor(scale_num / 4)), 128)
opt.min_nfc = min(opt.min_nfc_init * pow(2, math.floor(scale_num / 4)), 128)
opt.out_ = functions.generate_dir2save(opt)
opt.outf = '%s/%d' % (opt.out_,scale_num)
try:
os.makedirs(opt.outf)
except OSError:
pass
D_curr,G_curr, D_curr2,G_curr2 = init_models(opt)
if (nfc_prev==opt.nfc):
G_curr.load_state_dict(torch.load('%s/%d/netG.pth' % (opt.out_,scale_num-1)))
D_curr.load_state_dict(torch.load('%s/%d/netD.pth' % (opt.out_,scale_num-1)))
G_curr2.load_state_dict(torch.load('%s/%d/netG2.pth' % (opt.out_,scale_num-1)))
D_curr2.load_state_dict(torch.load('%s/%d/netD2.pth' % (opt.out_,scale_num-1)))
z_curr,in_s,G_curr, z_curr2,in_s2,G_curr2 = train_single_scale(D_curr,G_curr, reals,Gs,Zs,in_s,NoiseAmp, errD_plot,errG_plot,rec_loss_plot,cyc_loss_plot, D_curr2,G_curr2, reals2,Gs2,Zs2,in_s2,NoiseAmp2, errD2_plot,errG2_plot,rec_loss2_plot,cyc_loss2_plot, opt,scale_num)
G_curr = functions.reset_grads(G_curr,False)
G_curr.eval()
D_curr = functions.reset_grads(D_curr,False)
D_curr.eval()
G_curr2 = functions.reset_grads(G_curr2,False)
G_curr2.eval()
D_curr2 = functions.reset_grads(D_curr2,False)
D_curr2.eval()
Gs.append(G_curr)
Zs.append(z_curr)
NoiseAmp.append(opt.noise_amp)
Gs2.append(G_curr2)
Zs2.append(z_curr2)
NoiseAmp2.append(opt.noise_amp2)
torch.save(Zs, '%s/Zs.pth' % (opt.out_))
torch.save(Gs, '%s/Gs.pth' % (opt.out_))
torch.save(reals, '%s/reals.pth' % (opt.out_))
torch.save(NoiseAmp, '%s/NoiseAmp.pth' % (opt.out_))
torch.save(Zs2, '%s/Zs2.pth' % (opt.out_))
torch.save(Gs2, '%s/Gs2.pth' % (opt.out_))
torch.save(reals2, '%s/reals2.pth' % (opt.out_))
torch.save(NoiseAmp2, '%s/NoiseAmp2.pth' % (opt.out_))
scale_num+=1
nfc_prev = opt.nfc
del D_curr,G_curr, D_curr2,G_curr2
functions.my_plot(errD_plot,errG_plot,rec_loss_plot,cyc_loss_plot,opt)
functions.my_plot2(errD2_plot,errG2_plot,rec_loss2_plot,cyc_loss2_plot,opt)
return
parser.add_argument('--input_dir', help='input image dir', required=True)
parser.add_argument('--input_name', help='input image name', required=True)
parser.add_argument('--mode', help='task to be done', default='transfer')
parser.add_argument('--start_scale',
help='injection scale',
type=int,
default='0')
opt = parser.parse_args()
opt = functions.post_config(opt)
Gs = []
Zs = []
reals = []
NoiseAmp = []
Gs2 = []
dir2save = functions.generate_dir2save(opt)
if dir2save is None:
print('task does not exist')
else:
try:
os.makedirs(dir2save)
except OSError:
pass
real_in = functions.read_image(opt)
functions.adjust_scales2image(real_in, opt)
real_ = functions.read_image(opt)
real = imresize(real_, opt.scale1, opt)
reals = functions.creat_reals_pyramid(real, reals, opt)
Gs, Zs, NoiseAmp, Gs2 = functions.load_model(opt)
TuiGAN_transfer(Gs, Zs, reals, NoiseAmp, Gs2, opt)
ematrix = [0.0,0.0,0.0,0.0]
nframe = 0
for path in vpaths:
cap = cv2.VideoCapture(path)
while 1:
try:
frame = cv2.resize(cap.read()[1], (cfg.image_width,cfg.image_height))
original = copy.deepcopy(frame)
centre = frame[frame.shape[0]//4:3*frame.shape[0]//4,frame.shape[1]//4:3*frame.shape[1]//4,:] if cfg.method in SR_METHODS else None
if cfg.method == 'fsrcnn':
frame = np.array(Image.fromarray(frame).resize((cfg.image_width//4,cfg.image_height//4), resample=pil_image.BICUBIC))
elif cfg.method in ['vdsr','edsr','carn']:
frame = imresize(frame,scalar_scale=0.25)
elif cfg.method in ['sr']:
frame = cv2.resize(frame, (cfg.image_width//4,cfg.image_height//4),interpolation=cv2.INTER_LINEAR)
else:
frame = frame
centre = cv2.imencode(".jpg", centre, [cv2.IMWRITE_JPEG_QUALITY, cfg.jpg_quality])[1] if cfg.method in SR_METHODS else None # convert to jpg
centre = cv2.imdecode(centre, cv2.IMREAD_COLOR) if cfg.method in SR_METHODS else None
frame = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, cfg.jpg_quality])[1] # convert to jpg
frame = cv2.imdecode(frame, cv2.IMREAD_COLOR)
if cfg.method in SR_METHODS:
lr = frame
# preprocess
rms = np.sqrt(np.mean(np.power(diff, 2)))
print("RMSE: {:.4f}".format(rms))
return rms
if __name__ == "__main__":
# Siberia, same area of Fig. 8 in the paper
print("Siberia")
im10, im20, imGT = readh5("S2B_MSIL1C_20170725_T43WFQ.mat", imGT=True)
SR20 = dsen2_20(im10, im20)
# Evaluation against the ground truth on the 20m resolution bands (simulated)
print("DSen2:")
RMSE(SR20, imGT)
print("Bicubic:")
RMSE(imresize(im20, 2), imGT)
fig = plt.figure(1)
ax = fig.add_subplot(111)
cax = ax.imshow(SR20[:, :, 2])
fig.colorbar(cax)
ax.set_title("Super-resolved band B6")
fig = plt.figure(2)
ax = fig.add_subplot(111)
cax = plt.imshow(np.abs(SR20[:, :, 4] - imGT[:, :, 4]), vmin=0, vmax=200)
fig.colorbar(cax)
ax.set_title("Absolute differences to the GT, band B11")
plt.show(block=False)
#
def predict(self, data, **kwargs):
# eng = meng.start_matlab()
# for name, param in self.model.named_parameters():
# a = param.clone().cpu().data.numpy()
# print(name, a.max(), a.min())
# print('\n')
save_dir = os.path.join(self.config['preds_dir'], kwargs['testset'])
mkdir_if_not_exist(save_dir)
self.model.eval()
psnr_list = []
ssim_list = []
b_psnr_list = []
b_ssim_list = []
with torch.no_grad():
for img_bundle in data:
# print(img_bundle['name'])
if "color" in self.config.keys() and self.config["color"]:
x = img_bundle['origin']
y = img_bundle['y']
multichannel = True
else:
x = img_bundle['x']
y = img_bundle['y']
(rows, cols, channel) = y.shape
y, _, _ = np.split(y, indices_or_sections=channel, axis=2)
multichannel = False
x = torch.from_numpy(x).float().view(1, -1, x.shape[0],
x.shape[1])
if self.config['cuda']:
x = x.cuda()
# print(x[:5])
lr_size = (x.shape[2], x.shape[3])
hr_size = img_bundle['size']
if self.config['progressive']:
inter_sizes = np_utils.interval_size(
lr_size, hr_size, self.config['max_gradual_scale'])
else:
inter_sizes = []
inter_sizes.append(hr_size)
if self.config['net'] == 'wmcnn':
preds = self.model(x)
preds = [p.data.cpu().numpy() for p in preds]
# preds = [matlab.double(p.data.cpu().numpy().squeeze().tolist()) for p in preds]
# preds = eng.idwt2(*preds, 'bior1.1')
preds = pywt.idwt2((preds[0], (preds[1:])), 'bior1.1')
else:
preds = self.model(x)
if isinstance(preds, list):
# Y-channel's pixels are within [16, 235]
preds = np.clip(preds[-1].data.cpu().numpy(), 16 / 255,
235 / 255).astype(np.float64)
# preds = np.clip(preds[-1].data.cpu().numpy(), 0, 1).astype(np.float64)
else:
try:
preds = preds.data.cpu().numpy()
except AttributeError:
preds = preds
# preds = preds.mul(255).clamp(0, 255).round().div(255)
preds = np.clip(preds, 16 / 255,
235 / 255).astype(np.float64)
# preds = np.clip(preds, 0, 1).astype(np.float64)
preds = preds.squeeze()
if len(preds.shape) == 3:
preds = preds.transpose([1, 2, 0])
preds = modcrop(preds.squeeze(), kwargs['upscale'])
preds_bd = shave(preds.squeeze(), kwargs['upscale'])
y = modcrop(y.squeeze(), kwargs['upscale'])
y_bd = shave(y.squeeze(), kwargs['upscale'])
# print(preds_bd.shape, y_bd.shape)
x = x.data.cpu().numpy().squeeze()
bic = imresize.imresize(x, scalar_scale=kwargs['upscale'])
bic = np.clip(bic, 16 / 255, 235 / 255).astype(np.float64)
bic = shave(bic.squeeze(), kwargs['upscale'])
b_psnr = measure.compare_psnr(bic, y_bd)
b_ssim = measure.compare_ssim(bic, y_bd)
b_psnr_list.append(b_psnr)
b_ssim_list.append(b_ssim)
m_psnr = measure.compare_psnr(preds_bd, y_bd)
m_ssim = measure.compare_ssim(preds_bd,
y_bd,
multichannel=multichannel)
print('PSNR of image {} is {}'.format(img_bundle['name'],
m_psnr))
print('SSIM of image {} is {}'.format(img_bundle['name'],
m_ssim))
psnr_list.append(m_psnr)
ssim_list.append(m_ssim)
self.save_preds(save_dir, preds, img_bundle, True)
print('Averaged PSNR is {}'.format(np.mean(np.array(psnr_list))))
print('Averaged SSIM is {}'.format(np.mean(np.array(ssim_list))))
print('Averaged BIC PSNR is {}'.format(np.mean(np.array(b_psnr_list))))
print('Averaged BIC SSIM is {}'.format(np.mean(np.array(b_ssim_list))))
def predict(self, data, **kwargs):
# eng = meng.start_matlab()
# for name, param in self.model.named_parameters():
# a = param.clone().cpu().data.numpy()
# print(name, a.max(), a.min())
# print('\n')
cost_time = 0
save_dir = os.path.join(self.config['preds_dir'], kwargs['testset'])
mkdir_if_not_exist(save_dir)
self.model.eval()
psnr_list = []
ssim_list = []
b_psnr_list = []
b_ssim_list = []
gap_list = {}
diversity = 0
with torch.no_grad():
for img_bundle in data:
# print(img_bundle['name'])
if "color" in self.config.keys() and self.config["color"]:
x = img_bundle['origin']
y = img_bundle['y']
multichannel = True
else:
x = img_bundle['x']
y = img_bundle['y']
if len(y.shape) == 3:
(rows, cols, channel) = y.shape
y, _, _ = np.split(y, indices_or_sections=channel, axis=2)
else:
(rows, cols) = y.shape
multichannel = False
x = torch.from_numpy(x).float().view(1, -1, x.shape[0], x.shape[1])
if self.config['cuda']:
x = x.cuda()
# print(x[:5])
lr_size = (x.shape[2], x.shape[3])
hr_size = img_bundle['size']
if self.config['progressive']:
inter_sizes = np_utils.interval_size(lr_size, hr_size, self.config['max_gradual_scale'])
else:
inter_sizes = []
inter_sizes.append(hr_size)
start_time = time.time()
if self.config['net'] == 'rrgun':
preds = self.model(x, y_sizes=inter_sizes)
elif self.config['net'] == 'lapsrn':
# step = len(inter_sizes)
# if kwargs['upscale'] % 2 != 0:
# step = step + 1
step = int(np.ceil(math.log(kwargs['upscale'], 2)))
preds = self.model(x, step=step)[-1]
# y_numpy = preds[-1].data.cpu().numpy().squeeze()
# x = misc.imresize(y_numpy, size=hr_size,
# interp='bicubic', mode='F')
# x = np.array(x, dtype=np.float64)
# preds = torch.from_numpy(x)
# resize = tfs.Compose([tfs.ToPILImage(), tfs.Resize(hr_size, interpolation=Image.BICUBIC),
# tfs.ToTensor()])
# preds = resize(preds[-1].squeeze(0))
# preds = F.upsample(preds[-1], size=hr_size, mode='bilinear')
# preds = preds[-1]
elif self.config['net'] == 'lapgun':
preds = self.model(x, y_sizes=inter_sizes)
elif self.config['net'] in ['lapinternet', 'lapmtnet']:
# print(img.shape)
preds = self.model(x, size=inter_sizes[-1], step=self.config['step'])
elif self.config['net'] in ['ensemsr', 'stacksr', 'stacksr_back', 'stacksr_uni']:
input_list = self.em_generator(x)
preds, parts = self.model(input_list)
parts = torch.cat(tuple(parts), 0)
diversity += self.chisquare(parts)
# preds = com[-1]
elif self.config['net'] == 'wmcnn':
preds = self.model(x)
preds = [p.data.cpu().numpy() for p in preds]
# preds = [matlab.double(p.data.cpu().numpy().squeeze().tolist()) for p in preds]
# preds = eng.idwt2(*preds, 'bior1.1')
preds = pywt.idwt2((preds[0], (preds[1:])), 'bior1.1')
else:
preds = self.model(x)
# for c in com:
# c = c.data.cpu().numpy()
# continue
cost_time += time.time() - start_time
if isinstance(preds, list):
preds = np.clip(preds[-1].data.cpu().numpy(), 16/255, 235/255).astype(np.float64)
# preds = np.clip(preds[-1].data.cpu().numpy(), 0, 1).astype(np.float64)
else:
try:
preds = preds.data.cpu().numpy()
except AttributeError:
preds = preds
# preds = preds.mul(255).clamp(0, 255).round().div(255)
preds = np.clip(preds, 16/255, 235/255).astype(np.float64)
# preds = np.clip(preds, 0, 1).astype(np.float64)
preds = preds.squeeze()
if len(preds.shape) == 3:
preds = preds.transpose([1, 2, 0])
preds = modcrop(preds.squeeze(), kwargs['upscale'])
preds_bd = shave(preds.squeeze(), kwargs['upscale'])
y = modcrop(y.squeeze(), kwargs['upscale'])
# y = np.round(y * 255).astype(np.uint8)
y_bd = shave(y.squeeze(), kwargs['upscale'])#/ 255.
# print(preds_bd.shape, y_bd.shape)
x = x.data.cpu().numpy().squeeze()
# bic = x
bic = imresize.imresize(x, scalar_scale=kwargs['upscale'])
# bic = np.clip(bic, 16 / 255, 235 / 255).astype(np.float64)
bic = np.round(bic*255).astype(np.uint8)
bic = shave(bic.squeeze(), kwargs['upscale']) / 255.
b_psnr = measure.compare_psnr(bic, y_bd, data_range=1)
# b_ssim = measure.compare_ssim(bic, y_bd, data_range=1)
b_ssim = self.calculate_ssim(bic* 255, y_bd* 255)
# b_ssim = self.vifp_mscale(bic, y_bd)
b_psnr_list.append(b_psnr)
b_ssim_list.append(b_ssim)
m_psnr = measure.compare_psnr(preds_bd, y_bd)
# m_ssim = measure.compare_ssim(preds_bd, y_bd, multichannel=multichannel)
m_ssim = self.calculate_ssim(preds_bd* 255, y_bd* 255)
# print('image {} PSNR: {} SSIM: {}'.format(img_bundle['name'], m_psnr, m_ssim))
gap_list[m_psnr-b_psnr] = img_bundle['name']
psnr_list.append(m_psnr)
ssim_list.append(m_ssim)
test_value = '{}_{}'.format(m_psnr, m_ssim)
# self.save_preds(save_dir, test_value, preds, img_bundle, True)
diversity = diversity / len(data)
print('Averaged Diversity is {}'.format(diversity))
print('Averaged PSNR is {}, SSIM is {}'.format(np.mean(np.array(psnr_list)), np.mean(np.array(ssim_list))))
print('Averaged BIC PSNR is {}, SSIM is {}'.format(np.mean(np.array(b_psnr_list)), np.mean(np.array(b_ssim_list))))
# print(self.model.module.w_output)
# print(self.model.module.w_inter)
bigest_gap = sorted(gap_list, reverse=True)
print(bigest_gap)
print(gap_list[bigest_gap[0]], gap_list[bigest_gap[1]])
print('Inference cost time {}s'.format(cost_time))
