def _load_file(self, idx):
idx = self._get_index(idx)
lr_path = self.paths_LR[idx]
hr_path = self.paths_HR[idx]
lr = common.read_img(lr_path, self.opt['data_type'])
hr = common.read_img(hr_path, self.opt['data_type'])
return lr, hr, lr_path, hr_path
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = common.read_img(self.HR_env, HR_path)
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = common.read_img(self.LR_env, LR_path)
else:
raise NotImplementedError('Low resolution images do not exist')
img_LR, img_HR = self._get_patch(img_LR, img_HR, scale, HR_size, self.opt['phase'])
# channel conversion
if self.opt['color']:
img_LR, img_HR = common.channel_convert(img_HR.shape[2], [img_LR, img_HR], self.opt['color'])
# HWC to CHW, BGR to RGB, numpy to tensor
tensor_LR, tensor_HR = common.np2Tensor([img_LR, img_HR])
if LR_path is None:
LR_path = HR_path
return {'LR': tensor_LR, 'HR': tensor_HR, 'LR_path': LR_path, 'HR_path': HR_path}
def _load_file(self, idx):
lr_path = self.paths_LR[idx]
pan_path = self.paths_PAN[idx]
lr = common.read_img(lr_path, self.opt['data_type'])
pan = common.read_img(pan_path, self.opt['data_type'])
return lr, lr_path, pan, pan_path
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = common.read_img(self.HR_env, HR_path)
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = common.read_img(self.LR_env, LR_path)
else:
raise NotImplementedError('Low resolution images do not exist')
# get coeff
coeff_path = self.path_coeff[index]
coeff = common.read_coeff(self.coeff_env, coeff_path)
coeff = [coeff[0][i] for i in range(coeff.size)]
# channel conversion
if self.opt['color']:
img_LR, img_HR = common.channel_convert(img_HR.shape[2], [img_LR, img_HR], self.opt['color'])
# HWC to CHW, BGR to RGB, numpy to tensor
tensor_LR = torch.from_numpy(np.ascontiguousarray(img_LR)).float()
tensor_HR = torch.from_numpy(np.ascontiguousarray(img_HR)).float()
if LR_path is None:
LR_path = HR_path
return {'LR': tensor_LR, 'HR': tensor_HR, 'LR_path': LR_path, 'HR_path': HR_path, 'coeff': coeff}
def _load_file(self, idx):
idx = self._get_index(idx)
lr_path = self.paths_LR[idx]
lr = common.read_img(lr_path, self.opt['data_type'])
hr = []
hr_path = []
for i in range(len(self.scale)):
hr_path_ = self.paths_HR[i][idx]
hr_path.append(hr_path_)
hr.append(common.read_img(hr_path_, self.opt['data_type']))
return lr, hr, lr_path, hr_path
def _load_file(self, idx):
idx = self._get_index(idx)
ir_path = self.paths_IR[idx]
vis_path = self.paths_VIS[idx]
ir = common.read_img(self.IR_env, ir_path, self.opt['data_type'])
vis = common.read_img(self.VIS_env, vis_path, self.opt['data_type'])
if self.opt['phase'] == 'train' or self.opt['phase'] == 'val':
fus_path = self.paths_PF[idx]
fus = common.read_img(self.VIS_env, fus_path,
self.opt['data_type'])
return ir, vis, fus, ir_path, vis_path, fus_path
else:
return ir, vis, ir_path, vis_path
def __getitem__(self, index):
LR_path = None
# get LR image
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
H, W, C = img_LR.shape
# channel conversion
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
# HWC to CHW, BGR to RGB, numpy to tensor
if img_LR.shape[2] == 3:
img_LR = img_LR[:, :, [2, 1, 0]]
img_LR = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
return {'LR': img_LR, 'LR_path': LR_path}
def run(pretrained_path,
output_path,
model_name='SRFBN',
scale=4,
degrad='BI',
opt='options/test/test_SRFBN_example.json'):
opt = option.parse(opt)
opt = option.dict_to_nonedict(opt)
# model = create_model(opt)
model = define_net({
"scale": scale,
"which_model": "SRFBN",
"num_features": 64,
"in_channels": 3,
"out_channels": 3,
"num_steps": 4,
"num_groups": 6
})
img = common.read_img('./results/LR/MyImage/chip.png', 'img')
np_transpose = np.ascontiguousarray(img.transpose((2, 0, 1)))
tensor = torch.from_numpy(np_transpose).float()
lr_tensor = torch.unsqueeze(tensor, 0)
checkpoint = torch.load(pretrained_path)
if 'state_dict' in checkpoint.keys():
checkpoint = checkpoint['state_dict']
load_func = model.load_state_dict
load_func(checkpoint)
torch.save(model, './model.pt')
with torch.no_grad():
SR = model(lr_tensor)[0]
# visuals = np.transpose(SR.data[0].float().cpu().numpy(), (1, 2, 0)).astype(np.uint8)
visuals = np.transpose(SR.data[0].float().cpu().numpy(),
(1, 2, 0)).astype(np.uint8)
imageio.imwrite(output_path, visuals)
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# modcrop in validation phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, 8)
# get segmentation probability map
seg = torch.load(HR_path.replace('/img/', '/bicseg/').replace('.png', '.pth'))
seg = np.transpose(seg.numpy(), (1, 2, 0))
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = seg.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, HR_size)
W_s = _mod(W_s, random_scale, scale, HR_size)
img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
seg = cv2.resize(np.copy(seg), (W_s, H_s), interpolation=cv2.INTER_NEAREST)
H, W, _ = img_HR.shape
# using matlab imresize
img_LR = util.imresize_np(img_HR, 1 / scale, True)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
H, W, C = img_LR.shape
if self.opt['phase'] == 'train':
LR_size = HR_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
seg = seg[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# augmentation - flip, rotate
img_LR, img_HR, seg = util.augment([img_LR, img_HR, seg], self.opt['use_flip'],
self.opt['use_rot'])
# category
if 'building' in HR_path:
category = 0
elif 'plant' in HR_path:
category = 1
elif 'mountain' in HR_path:
category = 2
elif 'water' in HR_path:
category = 3
elif 'sky' in HR_path:
category = 4
elif 'grass' in HR_path:
category = 5
elif 'animal' in HR_path:
category = 6
else:
category = 7 # background
else:
category = -1 # during val, useless
# HWC to CHW, BGR to RGB, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
seg = torch.from_numpy(np.ascontiguousarray(np.transpose(seg, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {'LR': img_LR, 'HR': img_HR, 'seg': seg, 'category':category,
'LR_path': LR_path, 'HR_path': HR_path}
def get_noise_patch(self, patch_shape):
idx = random.randrange(0, self.n_noise_patches)
noise_patch_path = self.paths_noise_patches[idx]
noise_patch = common.read_img(noise_patch_path, self.opt['data_type'])
return noise_patch
