def test(self):
# networks
self.num_recursions = 16
self.model = Net(num_channels=self.num_channels,
base_filter=256,
num_recursions=self.num_recursions)
if self.gpu_mode:
self.model.cuda()
# load model
self.load_model()
# load dataset
test_data_loader = self.load_dataset(dataset='test')
# Test
print('Test is started.')
img_num = 0
self.model.eval()
for input, target in test_data_loader:
# input data (bicubic interpolated image)
if self.gpu_mode:
y_ = Variable(
utils.img_interp(input, self.scale_factor).cuda())
else:
y_ = Variable(utils.img_interp(input, self.scale_factor))
# prediction
_, recon_imgs = self.model(y_)
for i, recon_img in enumerate(recon_imgs):
img_num += 1
recon_img = recon_imgs[i].cpu().data
gt_img = target[i]
lr_img = input[i]
bc_img = utils.img_interp(input[i], self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
img_num,
save_dir=self.save_dir)
print("Saving %d test result images..." % img_num)
def test(self):
# networks
self.model = Net(num_channels=self.num_channels, base_filter=64, scale_factor=self.scale_factor)
if self.gpu_mode:
self.model.cuda()
# load model
self.load_model()
# load dataset
test_data_loader = self.load_dataset(dataset='test')
# Test
print('Test is started.')
img_num = 0
self.model.eval()
for input, target in test_data_loader:
# input data (low resolution image)
if self.gpu_mode:
y_ = Variable(input.cuda())
else:
y_ = Variable(input)
# prediction
recon_imgs = self.model(y_)
for i in range(self.test_batch_size):
img_num += 1
recon_img = recon_imgs[i].cpu().data
gt_img = utils.shave(target[i], border_size=8 * self.scale_factor)
lr_img = utils.shave(input[i], border_size=8)
bc_img = utils.shave(utils.img_interp(input[i], self.scale_factor), border_size=8 * self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs, psnrs, img_num, save_dir=self.save_dir)
print("Saving %d test result images..." % img_num)
def __init__(self, save_path, model_generator): self.path = save_path self.input = tf.placeholder(tf.float32,[None,input_image_size,input_image_size,3]) self.target = tf.placeholder(tf.float32,[None,output_image_size,output_image_size,3]) image_input = self.input - tf.reduce_mean(self.input) image_target = self.target - tf.reduce_mean(self.input) image_output = PeriodicShuffle(model_generator(image_input)) self.output = tf.clip_by_value(image_output + tf.reduce_mean(self.input), 0.0, 255.0) self.loss = tf.reduce_mean(tf.squared_difference(image_target, image_output)) self.PSNR = utils.PSNR(image_target, image_output) self.optimizer = tf.train.AdamOptimizer() self.session = tf.Session() self.saver = tf.train.Saver()
def train(self):
# networks
self.model = Net(num_channels=self.num_channels,
base_filter=64,
num_residuals=18)
# weigh initialization
self.model.weight_init()
# optimizer
self.momentum = 0.9
self.weight_decay = 0.0001
self.clip = 0.4
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay)
# loss function
if self.gpu_mode:
self.model.cuda()
self.MSE_loss = nn.MSELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset='train')
test_data_loader = self.load_dataset(dataset='test')
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.mkdir(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_input, test_target, test_bicubic = test_data_loader.dataset.__getitem__(
2)
test_input = test_input.unsqueeze(0)
test_target = test_target.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 10 every 20 epochs
if (epoch + 1) % 20 == 0:
for param_group in self.optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay: lr={}".format(
self.optimizer.param_groups[0]["lr"]))
epoch_loss = 0
for iter, (input, target, bi) in enumerate(train_data_loader):
# input data (bicubic interpolated image)
if self.gpu_mode:
x_ = Variable(target.cuda())
y_ = Variable(
utils.img_interp(input, self.scale_factor).cuda())
else:
x_ = Variable(target)
y_ = Variable(utils.img_interp(input, self.scale_factor))
# update network
self.optimizer.zero_grad()
recon_image = self.model(y_)
loss = self.MSE_loss(recon_image, x_)
loss.backward()
# gradient clipping
nn.utils.clip_grad_norm(self.model.parameters(), self.clip)
self.optimizer.step()
# log
epoch_loss += loss.item()
print("Epoch: [%2d] [%4d/%4d] loss: %.8f" %
((epoch + 1),
(iter + 1), len(train_data_loader), loss.item()))
# tensorboard logging
#logger.scalar_summary('loss', loss.item(), step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
recon_imgs = self.model(
Variable(
utils.img_interp(test_input, self.scale_factor).cuda()))
recon_img = recon_imgs[0].cpu().data
gt_img = test_target[0]
lr_img = test_input[0]
bc_img = utils.img_interp(test_input[0], self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
epoch + 1,
save_dir=self.save_dir,
is_training=True)
print("Saving training result images at epoch %d" % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
def predict():
random_date = "2020_02_03_22_18_02"
num_epochs = 399
test_path = "../data/test_img/"
img_list = ["man", "couple", "boat", "bridge", "cameraman", "house", "peppers"]
# test_path = "../data/Set12/"
# test_path = "../data/BSD68/"
hyp = {
"experiment_name": "poisson_ista15_stride7_peak1_11x11_169_onesided_tied_random_elu_smllr",
"dataset": "VOC",
"year": "2012",
"segmentation": False,
"image_set": "train",
"network": "DEA2Dtied",
"data_distribution": "poisson",
"model_distribution": "poisson",
"loss_distribution": "gaussian",
"dictionary_dim": 11,
"stride": 7,
"num_conv": 169,
"peak": 1,
"L": None,
"num_iters": 15,
"twosided": False,
"batch_size": 1,
"num_epochs": 400,
"zero_mean_filters": False,
"normalize": False,
"lr": 1e-4,
"lr_decay": 0.8,
"lr_step": 25,
"cyclic": False,
"amsgrad": False,
"info_period": 7000,
"lam": 0.01,
"nonlin": "ELU",
"single_bias": True,
"shuffle": True,
"crop_dim": (128, 128),
"init_with_DCT": False,
"init_with_saved_file": False,
"test_path": "../data/test_img/",
"denoising": True,
"mu": 0.0,
"supervised": True,
}
PATH = "../results/{}/{}".format(hyp["experiment_name"], random_date)
enable_plot = 0
T = 20
mu = 0
if test_path == "../data/test_img/":
num_test = len(img_list)
c = ["k", "b", "r", "g", "purple", "orange", "y", "m", "c"]
loss = torch.zeros(num_epochs, device=device)
psnr = np.zeros(num_epochs)
for epoch in range(num_epochs):
loss[epoch] = torch.tensor(
torch.load(os.path.join(PATH, "loss_epoch{}.pt".format(epoch)), device)
)
psnr[epoch] = np.load(os.path.join(PATH, "psnr_epoch{}.npy".format(epoch)))[0]
best_epoch = torch.argmin(loss)
print("best_epoch", best_epoch)
best_epoch = num_epochs
loss = loss.detach().cpu().numpy()
plt.figure(), plt.plot(loss), plt.xlabel("epochs"), plt.ylabel("loss")
plt.savefig(os.path.join(PATH, "loss.png"))
plt.figure(), plt.plot(psnr), plt.xlabel("epochs"), plt.ylabel("psnr")
plt.savefig(os.path.join(PATH, "psnr.png"))
print(hyp["experiment_name"])
print(random_date)
print("load data.")
test_loader = generator.get_path_loader(1, test_path, shuffle=False)
print("create model.")
net_init = torch.load(os.path.join(PATH, "model_init.pt"))
net = torch.load(os.path.join(PATH, "model_epoch{}.pt".format(best_epoch)))
if hyp["data_distribution"] == "binomial":
Jg = hyp["num_trials"]
with torch.no_grad():
ctr = 0
PSNR_noisy_list_all = list()
PSNR_init_list_all = list()
PSNR_list_all = list()
for img_org, tar in test_loader:
PSNR_noisy_list = list()
PSNR_init_list = list()
PSNR_list = list()
for t in range(T):
if hyp["data_distribution"] == "gaussian":
img_noisy = (
img_org + hyp["noiseSTD"] / 255 * torch.randn(img_org.shape)
).to(device)
img = img_org.to(device)
elif hyp["data_distribution"] == "binomial":
img = img_org.to(device)
sampler = torch.distributions.bernoulli.Bernoulli(probs=img)
img_noisy = sampler.sample()
for j in range(Jg - 1):
img_noisy += sampler.sample()
img_noisy /= Jg
elif hyp["data_distribution"] == "poisson":
img = img_org.to(device)
Q = torch.max(img) / hyp["peak"]
rate = img / Q
if torch.isnan(torch.min(rate)):
continue
sampler = torch.distributions.poisson.Poisson(rate)
if hyp["model_distribution"] == "poisson":
img_noisy = sampler.sample()
else:
img_noisy = sampler.sample() * Q
Hx_hat, _, _ = net_init(img_noisy, mu)
if hyp["model_distribution"] == "gaussian":
img_init_hat = Hx_hat + mu
elif hyp["model_distribution"] == "binomial":
img_init_hat = torch.nn.Sigmoid()(Hx_hat + mu)
elif hyp["model_distribution"] == "poisson":
img_init_hat = torch.exp(Hx_hat + mu) * Q
img_init_hat = img_init_hat[0, 0, :, :].detach().cpu().numpy()
Hx_hat, _, _ = net(img_noisy, mu)
if hyp["model_distribution"] == "gaussian":
img_hat = Hx_hat + mu
elif hyp["model_distribution"] == "binomial":
img_hat = torch.nn.Sigmoid()(Hx_hat + mu)
elif hyp["model_distribution"] == "poisson":
img_hat = torch.exp(Hx_hat + mu) * Q
else:
print("data distribution not exist!")
if t == 0:
if test_path == "../data/test_img/":
torchvision.utils.save_image(
img,
os.path.join(PATH, "{}_clean.png".format(img_list[ctr])),
)
torchvision.utils.save_image(
img_noisy * Q,
os.path.join(PATH, "{}_noisy.png".format(img_list[ctr])),
)
torchvision.utils.save_image(
img_hat,
os.path.join(PATH, "{}_est.png".format(img_list[ctr])),
)
img_hat = img_hat[0, 0, :, :].detach().cpu().numpy()
img_noisy = img_noisy[0, 0, :, :].detach().cpu().numpy()
img = img[0, 0, :, :].detach().cpu().numpy()
PSNR_noisy = utils.PSNR(img, img_noisy)
PSNR_init = utils.PSNR(img, img_init_hat)
PSNR = utils.PSNR(img, img_hat)
PSNR_noisy_list.append(PSNR_noisy)
PSNR_init_list.append(PSNR_init)
PSNR_list.append(PSNR)
PSNR_noisy = np.mean(PSNR_noisy_list)
PSNR_init = np.mean(PSNR_init_list)
PSNR = np.mean(PSNR_list)
PSNR_noisy_list_all.append(PSNR_noisy)
PSNR_init_list_all.append(PSNR_init)
PSNR_list_all.append(PSNR)
if test_path == "../data/test_img/":
print(img_list[ctr], PSNR_noisy, PSNR_init, PSNR)
ctr += 1
PSNR_noisy_all = np.mean(PSNR_noisy_list_all)
PSNR_init_all = np.mean(PSNR_init_list_all)
PSNR_all = np.mean(PSNR_list_all)
print(test_path, PSNR_noisy_all, PSNR_init_all, PSNR_all)
def main():
# ==========
# Load pre-trained model
# ==========
opts_dict = {
'radius': 3,
'stdf': {
'in_nc': 1,
'out_nc': 64,
'nf': 32,
'nb': 3,
'base_ks': 3,
'deform_ks': 3,
},
'qenet': {
'in_nc': 64,
'out_nc': 1,
'nf': 48,
'nb': 8,
'base_ks': 3,
},
}
model = MFVQE(opts_dict=opts_dict)
msg = f'loading model {ckp_path}...'
print(msg)
checkpoint = torch.load(ckp_path)
if 'module.' in list(checkpoint['state_dict'].keys())[0]: # multi-gpu training
new_state_dict = OrderedDict()
for k, v in checkpoint['state_dict'].items():
name = k[7:] # remove module
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
else: # single-gpu training
model.load_state_dict(checkpoint['state_dict'])
msg = f'> model {ckp_path} loaded.'
print(msg)
model = model.cuda()
model.eval()
# ==========
# Load entire video
# ==========
msg = f'loading raw and low-quality yuv...'
print(msg)
raw_y = utils.import_yuv(
seq_path=raw_yuv_path, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=True
)
lq_y, lq_u, lq_v = utils.import_yuv(
seq_path=lq_yuv_path, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=False
)
raw_y = raw_y.astype(np.float32) / 255.
lq_y = lq_y.astype(np.float32) / 255.
msg = '> yuv loaded.'
print(msg)
# ==========
# Define criterion
# ==========
criterion = utils.PSNR()
unit = 'dB'
# ==========
# Test
# ==========
pbar = tqdm(total=nfs, ncols=80)
ori_psnr_counter = utils.Counter()
enh_psnr_counter = utils.Counter()
for idx in range(nfs):
# load lq
idx_list = list(range(idx-3,idx+4))
idx_list = np.clip(idx_list, 0, nfs-1)
input_data = []
for idx_ in idx_list:
input_data.append(lq_y[idx_])
input_data = torch.from_numpy(np.array(input_data))
input_data = torch.unsqueeze(input_data, 0).cuda()
# enhance
enhanced_frm = model(input_data)
# eval
gt_frm = torch.from_numpy(raw_y[idx]).cuda()
batch_ori = criterion(input_data[0, 3, ...], gt_frm)
batch_perf = criterion(enhanced_frm[0, 0, ...], gt_frm)
ori_psnr_counter.accum(volume=batch_ori)
enh_psnr_counter.accum(volume=batch_perf)
# display
pbar.set_description(
"[{:.3f}] {:s} -> [{:.3f}] {:s}"
.format(batch_ori, unit, batch_perf, unit)
)
pbar.update()
pbar.close()
ori_ = ori_psnr_counter.get_ave()
enh_ = enh_psnr_counter.get_ave()
print('ave ori [{:.3f}] {:s}, enh [{:.3f}] {:s}, delta [{:.3f}] {:s}'.format(
ori_, unit, enh_, unit, (enh_ - ori_) , unit
))
print('> done.')
pixel_end_d = (count_d + 1) * batch_generate_size
pixel_start_h = count_h * batch_generate_size
pixel_end_h = (count_h + 1) * batch_generate_size
pixel_start_w = count_w * batch_generate_size
pixel_end_w = (count_w + 1) * batch_generate_size
sub_block = utils.read_imagecollection(
file_path=default_path + str(count_d) + str(count_h) +
str(count_w),
image_format='bmp')
# os.remove('3D_VDSR_'+str(count_d)+str(count_h)+str(count_w))
#print ('current path is:'+'3D_VDSR_'+str(count_d)+str(count_h)+str(count_w),'shape of sub_block is:',sub_block.shape)
reconstruction_output[pixel_start_d:pixel_end_d,
pixel_start_h:pixel_end_h,
pixel_start_w:pixel_end_w] = sub_block
dataset_ori = utils.read_imagecollection(opt.oripath)
dataset_interp = utils.read_imagecollection(
opt.interpath, image_format='bmp') #读取interp path 路径,文件后缀要注意bmp或是jpg！！
#print ('======>Read original image from : ',opt.oripath,' Read low resolution images from : ',opt.interpath)
#print ('PSNR of interp:',utils.PSNR(dataset_interp[:reconstruction_size,:reconstruction_size,:reconstruction_size],dataset_ori[:reconstruction_size,:reconstruction_size,:reconstruction_size]))
# print ('PSNR of reconstructor:',utils.PSNR(reconstruction_output,dataset_ori[:reconstruction_size,:reconstruction_size,:reconstruction_size]))
print(
utils.PSNR(
reconstruction_output,
dataset_ori[:reconstruction_size, :reconstruction_size, :
reconstruction_size]))
# Image.fromarray(reconstruction_output[..., 100]).show()
# Image.fromarray(dataset_interp[..., 100]).show()
utils.generate_2Dimage(save_mode=os.path.join(default_path, sub_name),
array_like=reconstruction_output)
def main():
# ==========
# parameters
# ==========
opts_dict = receive_arg()
rank = opts_dict['train']['rank']
unit = opts_dict['train']['criterion']['unit']
num_iter = int(opts_dict['train']['num_iter'])
interval_print = int(opts_dict['train']['interval_print'])
interval_val = int(opts_dict['train']['interval_val'])
# ==========
# init distributed training
# ==========
if opts_dict['train']['is_dist']:
utils.init_dist(
local_rank=rank,
backend='nccl'
)
# TO-DO: load resume states if exists
pass
# ==========
# create logger
# ==========
if rank == 0:
log_dir = op.join("exp", opts_dict['train']['exp_name'])
utils.mkdir(log_dir)
log_fp = open(opts_dict['train']['log_path'], 'w')
# log all parameters
msg = (
f"{'<' * 10} Hello {'>' * 10}\n"
f"Timestamp: [{utils.get_timestr()}]\n"
f"\n{'<' * 10} Options {'>' * 10}\n"
f"{utils.dict2str(opts_dict)}"
)
print(msg)
log_fp.write(msg + '\n')
log_fp.flush()
# ==========
# TO-DO: init tensorboard
# ==========
pass
# ==========
# fix random seed
# ==========
seed = opts_dict['train']['random_seed']
# >I don't know why should rs + rank
utils.set_random_seed(seed + rank)
# ==========
# Ensure reproducibility or Speed up
# ==========
#torch.backends.cudnn.benchmark = False # if reproduce
#torch.backends.cudnn.deterministic = True # if reproduce
torch.backends.cudnn.benchmark = True # speed up
# ==========
# create train and val data prefetchers
# ==========
# create datasets
train_ds_type = opts_dict['dataset']['train']['type']
val_ds_type = opts_dict['dataset']['val']['type']
radius = opts_dict['network']['radius']
assert train_ds_type in dataset.__all__, \
"Not implemented!"
assert val_ds_type in dataset.__all__, \
"Not implemented!"
train_ds_cls = getattr(dataset, train_ds_type)
val_ds_cls = getattr(dataset, val_ds_type)
train_ds = train_ds_cls(
opts_dict=opts_dict['dataset']['train'],
radius=radius
)
val_ds = val_ds_cls(
opts_dict=opts_dict['dataset']['val'],
radius=radius
)
# create datasamplers
train_sampler = utils.DistSampler(
dataset=train_ds,
num_replicas=opts_dict['train']['num_gpu'],
rank=rank,
ratio=opts_dict['dataset']['train']['enlarge_ratio']
)
val_sampler = None # no need to sample val data
# create dataloaders
train_loader = utils.create_dataloader(
dataset=train_ds,
opts_dict=opts_dict,
sampler=train_sampler,
phase='train',
seed=opts_dict['train']['random_seed']
)
val_loader = utils.create_dataloader(
dataset=val_ds,
opts_dict=opts_dict,
sampler=val_sampler,
phase='val'
)
assert train_loader is not None
batch_size = opts_dict['dataset']['train']['batch_size_per_gpu'] * \
opts_dict['train']['num_gpu'] # divided by all GPUs
num_iter_per_epoch = math.ceil(len(train_ds) * \
opts_dict['dataset']['train']['enlarge_ratio'] / batch_size)
num_epoch = math.ceil(num_iter / num_iter_per_epoch)
val_num = len(val_ds)
# create dataloader prefetchers
tra_prefetcher = utils.CPUPrefetcher(train_loader)
val_prefetcher = utils.CPUPrefetcher(val_loader)
# ==========
# create model
# ==========
model = MFVQE(opts_dict=opts_dict['network'])
model = model.to(rank)
if opts_dict['train']['is_dist']:
model = DDP(model, device_ids=[rank])
"""
# load pre-trained generator
ckp_path = opts_dict['network']['stdf']['load_path']
checkpoint = torch.load(ckp_path)
state_dict = checkpoint['state_dict']
if ('module.' in list(state_dict.keys())[0]) and (not opts_dict['train']['is_dist']): # multi-gpu pre-trained -> single-gpu training
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove module
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print(f'loaded from {ckp_path}')
elif ('module.' not in list(state_dict.keys())[0]) and (opts_dict['train']['is_dist']): # single-gpu pre-trained -> multi-gpu training
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = 'module.' + k # add module
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print(f'loaded from {ckp_path}')
else: # the same way of training
model.load_state_dict(state_dict)
print(f'loaded from {ckp_path}')
"""
# ==========
# define loss func & optimizer & scheduler & scheduler & criterion
# ==========
# define loss func
assert opts_dict['train']['loss'].pop('type') == 'CharbonnierLoss', \
"Not implemented."
loss_func = utils.CharbonnierLoss(**opts_dict['train']['loss'])
# define optimizer
assert opts_dict['train']['optim'].pop('type') == 'Adam', \
"Not implemented."
optimizer = optim.Adam(
model.parameters(),
**opts_dict['train']['optim']
)
# define scheduler
if opts_dict['train']['scheduler']['is_on']:
assert opts_dict['train']['scheduler'].pop('type') == \
'CosineAnnealingRestartLR', "Not implemented."
del opts_dict['train']['scheduler']['is_on']
scheduler = utils.CosineAnnealingRestartLR(
optimizer,
**opts_dict['train']['scheduler']
)
opts_dict['train']['scheduler']['is_on'] = True
# define criterion
assert opts_dict['train']['criterion'].pop('type') == \
'PSNR', "Not implemented."
criterion = utils.PSNR()
#
start_iter = 0 # should be restored
start_epoch = start_iter // num_iter_per_epoch
# display and log
if rank == 0:
msg = (
f"\n{'<' * 10} Dataloader {'>' * 10}\n"
f"total iters: [{num_iter}]\n"
f"total epochs: [{num_epoch}]\n"
f"iter per epoch: [{num_iter_per_epoch}]\n"
f"val sequence: [{val_num}]\n"
f"start from iter: [{start_iter}]\n"
f"start from epoch: [{start_epoch}]"
)
print(msg)
log_fp.write(msg + '\n')
log_fp.flush()
# ==========
# evaluate original performance, e.g., PSNR before enhancement
# ==========
vid_num = val_ds.get_vid_num()
if opts_dict['train']['pre-val'] and rank == 0:
msg = f"\n{'<' * 10} Pre-evaluation {'>' * 10}"
print(msg)
log_fp.write(msg + '\n')
per_aver_dict = {}
for i in range(vid_num):
per_aver_dict[i] = utils.Counter()
pbar = tqdm(
total=val_num,
ncols=opts_dict['train']['pbar_len']
)
# fetch the first batch
val_prefetcher.reset()
val_data = val_prefetcher.next()
while val_data is not None:
# get data
gt_data = val_data['gt'].to(rank) # (B [RGB] H W)
lq_data = val_data['lq'].to(rank) # (B T [RGB] H W)
index_vid = val_data['index_vid'].item()
name_vid = val_data['name_vid'][0] # bs must be 1!
b, _, _, _, _ = lq_data.shape
# eval
batch_perf = np.mean(
[criterion(lq_data[i,radius,...], gt_data[i]) for i in range(b)]
) # bs must be 1!
# log
per_aver_dict[index_vid].accum(volume=batch_perf)
# display
pbar.set_description(
"{:s}: [{:.3f}] {:s}".format(name_vid, batch_perf, unit)
)
pbar.update()
# fetch next batch
val_data = val_prefetcher.next()
pbar.close()
# log
ave_performance = np.mean([
per_aver_dict[index_vid].get_ave() for index_vid in range(vid_num)
])
msg = "> ori performance: [{:.3f}] {:s}".format(ave_performance, unit)
print(msg)
log_fp.write(msg + '\n')
log_fp.flush()
if opts_dict['train']['is_dist']:
torch.distributed.barrier() # all processes wait for ending
if rank == 0:
msg = f"\n{'<' * 10} Training {'>' * 10}"
print(msg)
log_fp.write(msg + '\n')
# create timer
total_timer = utils.Timer() # total tra + val time of each epoch
# ==========
# start training + validation (test)
# ==========
model.train()
num_iter_accum = start_iter
for current_epoch in range(start_epoch, num_epoch + 1):
# shuffle distributed subsamplers before each epoch
if opts_dict['train']['is_dist']:
train_sampler.set_epoch(current_epoch)
# fetch the first batch
tra_prefetcher.reset()
train_data = tra_prefetcher.next()
# train this epoch
while train_data is not None:
# over sign
num_iter_accum += 1
if num_iter_accum > num_iter:
break
# get data
gt_data = train_data['gt'].to(rank) # (B [RGB] H W)
lq_data = train_data['lq'].to(rank) # (B T [RGB] H W)
b, _, c, _, _ = lq_data.shape
input_data = torch.cat(
[lq_data[:,:,i,...] for i in range(c)],
dim=1
) # B [R1 ... R7 G1 ... G7 B1 ... B7] H W
enhanced_data = model(input_data)
# get loss
optimizer.zero_grad() # zero grad
loss = torch.mean(torch.stack(
[loss_func(enhanced_data[i], gt_data[i]) for i in range(b)]
)) # cal loss
loss.backward() # cal grad
optimizer.step() # update parameters
# update learning rate
if opts_dict['train']['scheduler']['is_on']:
scheduler.step() # should after optimizer.step()
if (num_iter_accum % interval_print == 0) and (rank == 0):
# display & log
lr = optimizer.param_groups[0]['lr']
loss_item = loss.item()
msg = (
f"iter: [{num_iter_accum}]/{num_iter}, "
f"epoch: [{current_epoch}]/{num_epoch - 1}, "
"lr: [{:.3f}]x1e-4, loss: [{:.4f}]".format(
lr*1e4, loss_item
)
)
print(msg)
log_fp.write(msg + '\n')
if ((num_iter_accum % interval_val == 0) or \
(num_iter_accum == num_iter)) and (rank == 0):
# save model
checkpoint_save_path = (
f"{opts_dict['train']['checkpoint_save_path_pre']}"
f"{num_iter_accum}"
".pt"
)
state = {
'num_iter_accum': num_iter_accum,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
if opts_dict['train']['scheduler']['is_on']:
state['scheduler'] = scheduler.state_dict()
torch.save(state, checkpoint_save_path)
# validation
with torch.no_grad():
per_aver_dict = {}
for index_vid in range(vid_num):
per_aver_dict[index_vid] = utils.Counter()
pbar = tqdm(
total=val_num,
ncols=opts_dict['train']['pbar_len']
)
# train -> eval
model.eval()
# fetch the first batch
val_prefetcher.reset()
val_data = val_prefetcher.next()
while val_data is not None:
# get data
gt_data = val_data['gt'].to(rank) # (B [RGB] H W)
lq_data = val_data['lq'].to(rank) # (B T [RGB] H W)
index_vid = val_data['index_vid'].item()
name_vid = val_data['name_vid'][0] # bs must be 1!
b, _, c, _, _ = lq_data.shape
input_data = torch.cat(
[lq_data[:,:,i,...] for i in range(c)],
dim=1
) # B [R1 ... R7 G1 ... G7 B1 ... B7] H W
enhanced_data = model(input_data) # (B [RGB] H W)
# eval
batch_perf = np.mean(
[criterion(enhanced_data[i], gt_data[i]) for i in range(b)]
) # bs must be 1!
# display
pbar.set_description(
"{:s}: [{:.3f}] {:s}"
.format(name_vid, batch_perf, unit)
)
pbar.update()
# log
per_aver_dict[index_vid].accum(volume=batch_perf)
# fetch next batch
val_data = val_prefetcher.next()
# end of val
pbar.close()
# eval -> train
model.train()
# log
ave_per = np.mean([
per_aver_dict[index_vid].get_ave() for index_vid in range(vid_num)
])
msg = (
"> model saved at {:s}\n"
"> ave val per: [{:.3f}] {:s}"
).format(
checkpoint_save_path, ave_per, unit
)
print(msg)
log_fp.write(msg + '\n')
log_fp.flush()
if opts_dict['train']['is_dist']:
torch.distributed.barrier() # all processes wait for ending
# fetch next batch
train_data = tra_prefetcher.next()
# end of this epoch (training dataloader exhausted)
# end of all epochs
# ==========
# final log & close logger
# ==========
if rank == 0:
total_time = total_timer.get_interval() / 3600
msg = "TOTAL TIME: [{:.1f}] h".format(total_time)
print(msg)
log_fp.write(msg + '\n')
msg = (
f"\n{'<' * 10} Goodbye {'>' * 10}\n"
f"Timestamp: [{utils.get_timestr()}]"
)
print(msg)
log_fp.write(msg + '\n')
log_fp.close()
def run(model_name,
base_dir,
data_dir,
task='sc',
train_batch_size=128,
eval_batch_size=1024,
epochs=125,
mode='train'):
_NUM_TRAIN_IMAGES = FLAGS.num_train_images
_NUM_EVAL_IMAGES = FLAGS.num_val_images
if task == 'sc' or task == 'lasso':
training_steps_per_epoch = int(_NUM_TRAIN_IMAGES // train_batch_size)
validation_steps_per_epoch = int(_NUM_EVAL_IMAGES // eval_batch_size)
elif task == 'cs':
training_steps_per_epoch = 3125
validation_steps_per_epoch = 10
_BASE_LR = FLAGS.base_lr
# Deal with paths of data, checkpoints and logs
base_dir = os.path.abspath(base_dir)
model_dir = os.path.join(base_dir, 'models', FLAGS.exp_name,
'replicate_' + str(FLAGS.replicate))
log_dir = os.path.join(base_dir, 'logs', FLAGS.exp_name,
'replicate_' + str(FLAGS.replicate))
logging.info('Saving checkpoints at %s', model_dir)
logging.info('Saving tensorboard summaries at %s', log_dir)
logging.info('Use training batch size: %s.', train_batch_size)
logging.info('Use eval batch size: %s.', eval_batch_size)
logging.info('Training model using data_dir in directory: %s', data_dir)
if task == 'sc' or task == 'lasso':
A = np.load(os.path.join(data_dir, 'A.npy'),
allow_pickle=True).astype(np.float32)
M, N = A.shape
F = None
D = None
elif task == 'cs':
A = np.load(os.path.join(data_dir, 'A_128_512.npy'),
allow_pickle=True).astype(np.float32)
D = np.load(os.path.join(data_dir, 'D_256_512.npy'),
allow_pickle=True).astype(np.float32)
N = D.shape[0]
F = D.shape[1]
else:
raise ValueError('invalid task type')
if FLAGS.model_name.startswith('alista'):
alista_W = np.load(os.path.join(data_dir, 'W.npy'),
allow_pickle=True).astype(np.float32)
np.random.seed(FLAGS.seed)
if mode == 'train':
train_dataset = data_preprocessing.input_fn(True,
data_dir,
train_batch_size,
task,
drop_remainder=False,
A=A)
val_dataset = data_preprocessing.input_fn(False,
data_dir,
eval_batch_size,
task,
drop_remainder=False,
A=A)
summary_writer = tf.summary.create_file_writer(log_dir)
# Define a Lista model
if FLAGS.model_name == 'lista':
model = models.Lista(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.share_W,
D,
name='Lista')
output_interval = N
elif FLAGS.model_name == 'lfista':
model = models.Lfista(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.share_W,
D,
name='Lfista')
output_interval = N
elif FLAGS.model_name == 'lamp':
model = models.Lamp(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.share_W,
D,
name='Lamp')
output_interval = M + N
elif FLAGS.model_name == 'step_lista':
assert FLAGS.model_lam == FLAGS.lasso_lam
model = models.StepLista(A,
FLAGS.num_layers,
FLAGS.lasso_lam,
D,
name='StepLista')
output_interval = N
elif FLAGS.model_name == 'lista_cp':
model = models.ListaCp(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.share_W,
D,
name='ListaCp')
output_interval = N
elif FLAGS.model_name == 'lista_cpss':
model = models.ListaCpss(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.ss_q_per_layer,
FLAGS.ss_maxq,
FLAGS.share_W,
D,
name='ListaCpss')
output_interval = N
elif FLAGS.model_name == 'alista':
model = models.Alista(A,
alista_W,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.ss_q_per_layer,
FLAGS.ss_maxq,
D,
name='Alista')
output_interval = N
elif FLAGS.model_name == 'glista':
model = models.Glista(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.ss_q_per_layer,
FLAGS.ss_maxq,
FLAGS.share_W,
D,
name='Glista',
alti=FLAGS.glista_alti,
gain_func=FLAGS.gain_func)
output_interval = N * 2
elif FLAGS.model_name == 'tista':
assert FLAGS.tista_sigma2 is not None
model = models.Tista(A,
FLAGS.num_layers,
FLAGS.model_lam,
FLAGS.tista_sigma2,
FLAGS.share_W,
D,
name='Tista')
output_interval = N
else:
raise NotImplementedError(
'Other types of models not are not implemented yet')
var_list = {}
checkpoint = tf.train.Checkpoint(model=model)
prev_model, prev_layer = utils.check_and_load_partial(
model_dir, FLAGS.num_layers)
if task == 'lasso':
_A_const = tf.constant(A, name='A_lasso_const')
loss = utils.LassoLoss(_A_const, FLAGS.lasso_lam, N, F)
metrics_compile = [
utils.LassoObjective('lasso', _A_const, FLAGS.lasso_lam, M, N, -1)
]
monitor = 'val_lasso'
else:
loss = utils.MSE(N, F)
metrics_compile = [utils.NMSE('nmse', N)]
monitor = 'val_nmse'
if mode == 'test':
if prev_layer != FLAGS.num_layers:
raise ValueError('Should have a fully trained model!')
checkpoint.restore(prev_model).assert_existing_objects_matched()
res_dict = {}
eval_files = FLAGS.test_files
if task == 'lasso':
# do layer-wise testing
test_metrics = [
utils.LassoObjective('lasso_layer{}'.format(i), _A_const,
FLAGS.lasso_lam, M, N, i)
for i in range(FLAGS.num_layers)
]
else:
test_metrics = [
utils.EvalNMSE('nmse_layer{}'.format(i), M, N, output_interval,
i) for i in range(FLAGS.num_layers)
]
for layer_id in range(FLAGS.num_layers):
model.create_cell(layer_id)
for i in range(len(eval_files)):
val_ds = data_preprocessing.input_fn(False,
data_dir,
eval_batch_size,
drop_remainder=False,
A=A,
filename=eval_files[i])
logging.info('Compiling model.')
model.compile(optimizer=tf.keras.optimizers.Adam(_BASE_LR),
loss=loss,
metrics=test_metrics)
metrics = model.evaluate(x=val_ds, verbose=2)
if task == 'lasso':
output = model.predict(x=val_ds, verbose=2)
final_xh = output[:, -N:]
eval_file_basename = os.path.basename(
eval_files[i]).strip('.npy')
np.save(
os.path.join(model_dir,
eval_file_basename + '_final_output.npy'),
final_xh)
res_dict[eval_files[i]] = metrics[1:]
for k, v in res_dict.items():
logging.info('%s : %s', k, str(v))
return
for layer_id in range(FLAGS.num_layers):
logging.info('Building Lista Keras model.')
model.create_cell(layer_id)
# Deal with the variables that have been trained in previous layers
for name in var_list:
var_list[name] += 1
# Deal with the variables in the current layer
for v in model.layers[layer_id].trainable_variables:
if v.name not in var_list:
var_list[v.name] = 0
if layer_id == prev_layer - 1 and prev_model:
checkpoint.restore(prev_model).assert_existing_objects_matched()
logging.info('Checkpoint restored from %s.', prev_model)
model.compile(optimizer=tf.keras.optimizers.Adam(_BASE_LR),
loss=loss,
metrics=metrics_compile)
metrics = model.evaluate(x=val_dataset, verbose=2)
val_metric = metrics[1]
with summary_writer.as_default():
for value, key in zip(metrics, model.metrics_names):
tf.summary.scalar(key, value, layer_id + 1)
continue
elif layer_id < prev_layer:
logging.info('Skip layer %d.', layer_id + 1)
continue
logging.info('Compiling model.')
model.compile(optimizer=utils.Adam(var_list,
True,
learning_rate=_BASE_LR),
loss=loss,
metrics=metrics_compile)
earlystopping_cb = tf.keras.callbacks.EarlyStopping(
monitor=monitor,
min_delta=0,
patience=5,
mode='min',
restore_best_weights=False)
cbs = [earlystopping_cb]
logging.info('Fitting Lista Keras model.')
model.fit(train_dataset,
epochs=epochs,
steps_per_epoch=training_steps_per_epoch,
callbacks=cbs,
validation_data=val_dataset,
validation_steps=validation_steps_per_epoch,
verbose=2)
logging.info('Finished fitting Lista Keras model.')
model.summary()
for i in range(2):
logging.info('Compiling model.')
model.compile(optimizer=utils.Adam(var_list,
learning_rate=_BASE_LR * 0.2 *
0.1**i),
loss=loss,
metrics=metrics_compile)
earlystopping_cb = tf.keras.callbacks.EarlyStopping(
monitor=monitor,
min_delta=0,
patience=5,
mode='min',
restore_best_weights=False)
cbs = [earlystopping_cb]
logging.info('Fitting Lista Keras model.')
history = model.fit(train_dataset,
epochs=epochs,
steps_per_epoch=training_steps_per_epoch,
callbacks=cbs,
validation_data=val_dataset,
validation_steps=validation_steps_per_epoch,
verbose=2)
logging.info('Finished fitting Lista Keras model.')
model.summary()
val_metric = history.history[monitor][-1]
with summary_writer.as_default():
for key in history.history.keys():
tf.summary.scalar(key, history.history[key][-1], layer_id + 1)
try:
checkpoint.save(utils.save_partial(model_dir, layer_id))
logging.info('Checkpoint saved at %s',
utils.save_partial(model_dir, layer_id))
except tf.errors.NotFoundError:
pass
if task == 'cs':
raise NotImplementedError(
'Compressive sensing testing part not implemented yet')
data = np.load(os.path.join(data_dir, 'set11.npy'), allow_pickle=True)
phi = np.load(os.path.join(data_dir, 'phi_128_256.npy'),
allow_pickle=True).astype(np.float32)
psnr = utils.PSNR()
for im in data:
im_ = im.astype(np.float32)
cols = utils.im2cols(im_)
patch_mean = np.mean(cols, axis=1, keepdims=True)
fs = ((cols - patch_mean) / 255.0).astype(np.float32)
ys = np.matmul(fs, phi.transpose())
fs_rec = model.predict_on_batch(ys)[:, -N:]
cols_rec = fs_rec * 255.0 + patch_mean
im_rec = utils.col2im(cols_rec).astype(np.float32)
psnr.update_state(im_, im_rec)
logging.info('Test PSNR: %f', psnr.result().numpy())
val_metric = float(psnr.result().numpy())
with summary_writer.as_default():
tf.summary.scalar('test_psnr', psnr.result().numpy(), 0)
return val_metric
# utils.show_image_list(img_list_gauss_mean, "Image with gauss after mean filter!")
# utils.show_image_list(img_list_gauss_neavf, "Image with gauss after NEAVF filter!")
# Noise measure
psnr_gauss = []
psnr_impulsive = []
psnr_NEAVF_impulsive = []
psnr_NEAVF_gauss = []
mae_gauss = []
mae_impulsive = []
mae_NEAVF_impulsive = []
mae_NEAVF_gauss = []
img_list.append(img_list[0])
img_list.append(img_list[1])
for i in range(len(img_list)):
psnr_gauss.append(utils.PSNR(img_list[i], img_list_gauss_mean[i]))
psnr_impulsive.append(utils.PSNR(img_list[i],
img_list_impulsive_median[i]))
psnr_NEAVF_impulsive.append(
utils.PSNR(img_list[i], img_list_impulsive_neavf[i]))
psnr_NEAVF_gauss.append(utils.PSNR(img_list[i], img_list_gauss_neavf[i]))
print("PSNR for Image {0} after gaussian noise + filter(median) = {1};".
format(i, psnr_gauss[i]))
print("PSNR for Image {0} after impulsive noise + filter(median) = {1};".
format(i, psnr_impulsive[i]))
print("PSNR for Image {0} after impulsive noise + filter(NEAVF) = {1};".
format(i, psnr_NEAVF_impulsive[i]))
print("PSNR for Image {0} after gauss noise + filter(NEAVF) = {1};".format(
i, psnr_NEAVF_gauss[i]))
mae_gauss.append(utils.MAE(img_list[i], img_list_gauss[i]))
mae_impulsive.append(utils.MAE(img_list[i], img_list_impulsive[i]))
def test(self):
# networks
self.G = Generator(num_channels=self.num_channels,
base_filter=64,
num_residuals=16)
if self.gpu_mode:
self.G.cuda()
# load model
self.load_model()
# load dataset
for test_dataset in self.test_dataset:
test_data_loader = self.load_dataset(dataset=test_dataset,
is_train=False)
# Test
print('Test is started.')
img_num = 0
total_img_num = len(test_data_loader)
self.G.eval()
for lr, hr, bc in test_data_loader:
# input data (low resolution image)
if self.num_channels == 1:
y_ = Variable(utils.norm(lr[:, 0].unsqueeze(1), vgg=True))
else:
y_ = Variable(utils.norm(lr, vgg=True))
if self.gpu_mode:
y_ = y_.cuda()
# prediction
recon_imgs = self.G(y_)
for i, recon_img in enumerate(recon_imgs):
img_num += 1
sr_img = utils.denorm(recon_img.cpu().data, vgg=True)
# save result image
save_dir = os.path.join(self.save_dir, 'test_result',
test_dataset)
utils.save_img(sr_img, img_num, save_dir=save_dir)
# calculate psnrs
if self.num_channels == 1:
gt_img = hr[i][0].unsqueeze(0)
lr_img = lr[i][0].unsqueeze(0)
bc_img = bc[i][0].unsqueeze(0)
else:
gt_img = hr[i]
lr_img = lr[i]
bc_img = bc[i]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
img_num,
save_dir=save_dir)
print('Test DB: %s, Saving result images...[%d/%d]' %
(test_dataset, img_num, total_img_num))
print('Test is finishied.')
def test(ARGS):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
SCALE = ARGS['SCALE']
path = ARGS["TESTIMG"]
fullimg = cv2.imread(path, 3)
width = fullimg.shape[0]
height = fullimg.shape[1]
cropped = fullimg[0:(width - (width % SCALE)),
0:(height - (height % SCALE)), :]
img = cv2.resize(cropped,
None,
fx=1. / SCALE,
fy=1. / SCALE,
interpolation=cv2.INTER_CUBIC)
floatimg = img.astype(np.float32) / 255.0
# Convert to YCbCr color space
imgYCbCr = cv2.cvtColor(floatimg, cv2.COLOR_BGR2YCrCb)
imgY = imgYCbCr[:, :, 0]
LR_input_ = imgY.reshape(1, imgY.shape[0], imgY.shape[1], 1)
with tf.Session(config=config) as sess:
print("\nStart running tests on the model\n")
# #load the model with tf.data generator
ckpt_name = ARGS["CKPT"] + ".meta"
saver = tf.train.import_meta_graph(ckpt_name)
saver.restore(sess, tf.train.latest_checkpoint(ARGS["CKPT_dir"]))
graph_def = sess.graph
LR_tensor = sess.graph.get_tensor_by_name("IteratorGetNext:0")
inp = cv2.cvtColor(
(cropped.astype(np.float32) / 255.0),
cv2.COLOR_BGR2YCrCb)[:, :, 0].reshape(1, cropped.shape[0],
cropped.shape[1], 1)
bicub = cv2.cvtColor(
cv2.resize(cropped,
None,
fx=SCALE,
fy=SCALE,
interpolation=cv2.INTER_CUBIC), cv2.COLOR_BGR2YCrCb)
output = sess.run(sess.graph.get_tensor_by_name("NCHW_output:0"),
feed_dict={LR_tensor: inp})
Y = output[0][0]
cv2.imshow('LapSRN HR image', Y)
cv2.imshow('Bicubic HR image', bicub[:, :, 0])
cv2.waitKey(0)
LR_tensor = graph_def.get_tensor_by_name("IteratorGetNext:0")
HR_tensor = graph_def.get_tensor_by_name("NCHW_output:0")
output = sess.run(HR_tensor, feed_dict={LR_tensor: LR_input_})
Y = output[0][0]
Y = Y.reshape(Y.shape[0], Y.shape[1], 1)
Cr = np.expand_dims(cv2.resize(imgYCbCr[:, :, 1],
None,
fx=SCALE,
fy=SCALE,
interpolation=cv2.INTER_CUBIC),
axis=2)
Cb = np.expand_dims(cv2.resize(imgYCbCr[:, :, 2],
None,
fx=SCALE,
fy=SCALE,
interpolation=cv2.INTER_CUBIC),
axis=2)
HR_image_YCrCb = np.concatenate((Y, Cr, Cb), axis=2)
HR_image = ((cv2.cvtColor(HR_image_YCrCb, cv2.COLOR_YCrCb2BGR)) *
255.0).clip(min=0, max=255)
HR_image = (HR_image).astype(np.uint8)
bicubic_image = cv2.resize(img,
None,
fx=SCALE,
fy=SCALE,
interpolation=cv2.INTER_CUBIC)
print(np.amax(Y), np.amax(LR_input_))
print("PSNR of LapSRN generated image: ",
utils.PSNR(cropped, HR_image))
print("PSNR of bicubic interpolated image: ",
utils.PSNR(cropped, bicubic_image))
cv2.imshow('Original image', fullimg)
cv2.imshow('HR image', HR_image)
cv2.imshow('Bicubic HR image', bicubic_image)
cv2.waitKey(0)
def test(ARGS):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
SCALE = ARGS['SCALE']
path = ARGS["TESTIMG"]
fullimg = cv2.imread(path, 3)
width = fullimg.shape[0]
height = fullimg.shape[1]
cropped = fullimg[0:(width - (width % SCALE)), 0:(height - (height % SCALE)), :]
img = cv2.resize(cropped, None, fx=1. / SCALE, fy=1. / SCALE, interpolation=cv2.INTER_CUBIC)
floatimg = img.astype(np.float32) / 255.0
LR_input_ = floatimg.reshape(1, floatimg.shape[0], floatimg.shape[1], 3)
# with tf.gfile.GFile("nchw_frozen_ESPCN_graph_x2.pb", 'rb') as f:
# graph_def = tf.GraphDef()
# graph_def.ParseFromString(f.read())
# g = tf.import_graph_def(graph_def)
# # output_tensor = sess.graph.get_tensor_by_name("IteratorGetNext")
#
# sess = tf.Session(graph=g,config=config)
#
# LR_tensor = sess.graph.get_tensor_by_name("import/IteratorGetNext:0")
# inp = cv2.cvtColor((cropped.astype(np.float32) / 255.0), cv2.COLOR_BGR2YCrCb)[:,:,0].reshape(1, cropped.shape[0], cropped.shape[1], 1)
# output = sess.run(sess.graph.get_tensor_by_name("import/NCHW_output:0"), feed_dict={LR_tensor: inp})
# Y = output[0][0]
# print(Y.shape)
# cv2.imshow('Bicubic HR image', Y)
# cv2.waitKey(0)
with tf.Session(config=config) as sess:
print("\nStart running tests on the model\n")
# #load the model with tf.data generator
ckpt_name = ARGS["CKPT"] + ".meta"
saver = tf.train.import_meta_graph(ckpt_name)
saver.restore(sess, tf.train.latest_checkpoint(ARGS["CKPT_dir"]))
graph_def = sess.graph
LR_tensor = graph_def.get_tensor_by_name("IteratorGetNext:0")
HR_tensor = graph_def.get_tensor_by_name("NHWC_output:0")
output = sess.run(HR_tensor, feed_dict={LR_tensor: LR_input_})
Y = output[0]
'''
Cr = np.expand_dims(cv2.resize(imgYCbCr[:, :, 1], None, fx=SCALE, fy=SCALE, interpolation=cv2.INTER_CUBIC),
axis=2)
Cb = np.expand_dims(cv2.resize(imgYCbCr[:, :, 2], None, fx=SCALE, fy=SCALE, interpolation=cv2.INTER_CUBIC),
axis=2)
'''
# HR_image_YCrCb = np.concatenate((Y, Cr, Cb), axis=2)
HR_image = ((Y+1) * 127.5).clip(min=0, max=255)
HR_image = (HR_image).astype(np.uint8)
bicubic_image = cv2.resize(img, None, fx=SCALE, fy=SCALE, interpolation=cv2.INTER_CUBIC)
print(np.amax(Y), np.amax(LR_input_))
print("PSNR of ESPCN generated image: ", utils.PSNR(cropped, HR_image))
print("PSNR of bicubic interpolated image: ", utils.PSNR(cropped, bicubic_image))
cv2.imwrite('./singletest/Ori.png', cropped)
cv2.imwrite('./singletest/ESPCN.png', HR_image)
cv2.imwrite('./singletest/bicubic.png', bicubic_image)
def train(self):
# networks
self.num_recursions = 16
self.model = Net(num_channels=self.num_channels,
base_filter=256,
num_recursions=self.num_recursions)
# weigh initialization
self.model.weight_init()
# optimizer
self.momentum = 0.9
self.weight_decay = 0.0001
self.loss_alpha = 1.0
self.loss_alpha_zero_epoch = 25
self.loss_alpha_decay = self.loss_alpha / self.loss_alpha_zero_epoch
self.loss_beta = 0.001
# learnable parameters
param_groups = list(self.model.parameters())
param_groups = [{'params': param_groups}]
param_groups += [{'params': [self.model.w]}]
self.optimizer = optim.Adam(param_groups, lr=self.lr)
# loss function
if self.gpu_mode:
self.model.cuda()
self.MSE_loss = nn.MSELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset='train')
test_data_loader = self.load_dataset(dataset='test')
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.mkdir(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_input, test_target = test_data_loader.dataset.__getitem__(2)
test_input = test_input.unsqueeze(0)
test_target = test_target.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 10 every 20 epochs
if (epoch + 1) % 20 == 0:
for param_group in self.optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay: lr={}".format(
self.optimizer.param_groups[0]["lr"]))
# loss_alpha decayed to zero after 25 epochs
self.loss_alpha = max(0.0, self.loss_alpha - self.loss_alpha_decay)
epoch_loss = 0
for iter, (input, target) in enumerate(train_data_loader):
# input data (bicubic interpolated image)
if self.gpu_mode:
y = Variable(target.cuda())
x = Variable(
utils.img_interp(input, self.scale_factor).cuda())
else:
y = Variable(target)
x = Variable(utils.img_interp(input, self.scale_factor))
# update network
self.optimizer.zero_grad()
y_d_, y_ = self.model(x)
# loss1
loss1 = 0
for d in range(self.num_recursions):
loss1 += (self.MSE_loss(y_d_[d], y) / self.num_recursions)
# loss2
loss2 = self.MSE_loss(y_, y)
# regularization
reg_term = 0
for theta in self.model.parameters():
reg_term += torch.mean(torch.sum(theta**2))
# total loss
loss = self.loss_alpha * loss1 + (
1 - self.loss_alpha) * loss2 + self.loss_beta * reg_term
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
print("Epoch: [%2d] [%4d/%4d] loss: %.8f" %
((epoch + 1),
(iter + 1), len(train_data_loader), loss.data[0]))
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
_, recon_imgs = self.model(
Variable(
utils.img_interp(test_input, self.scale_factor).cuda()))
recon_img = recon_imgs[0].cpu().data
gt_img = test_target[0]
lr_img = test_input[0]
bc_img = utils.img_interp(test_input[0], self.scale_factor)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
epoch + 1,
save_dir=self.save_dir,
is_training=True)
print("Saving training result images at epoch %d" % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
def train(self):
# networks
self.model = Net(3, 64, 3, 1)
LossNet.creat_loss_Net(self)
if self.gpu_mode:
self.model.cuda()
if self.pre_epochs == 0:
# weigh initialization
self.model.weight_init()
else:
# (self, is training or not, is E_model(True) or D_model(False))
utils.load_model(self, True, True)
# self.model.load_state_dict(torch.load(self.E_pretrain))
# optimizer
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
self.valid = None
self.fake = None
if 'A' in self.model_loss:
self.discriminator = Discriminator().cuda(
) if self.gpu_mode else Discriminator()
Tensor = torch.cuda.FloatTensor if self.gpu_mode else torch.Tensor
self.optimizer_D = torch.optim.Adam(
self.discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
if self.pre_epochs != 0:
utils.load_model(self, True, False)
if self.loss_F == "BCEWithLogitsLoss":
self.criterion_GAN = nn.BCEWithLogitsLoss(
size_average=False).cuda(
) if self.gpu_mode else nn.BCEWithLogitsLoss(
size_average=False)
# elif self.loss_F == "Cross"
# print('---------- Networks architecture -------------')
# utils.print_network(self.model)
# print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset=self.train_dataset,
is_train=True)
test_data_loader = self.load_dataset(dataset=self.test_dataset[0],
is_train=False)
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.makedirs(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
avg_loss_D = []
# test image
test_lr, test_hr, test_bc, name = test_data_loader.dataset.__getitem__(
2)
test_lr = test_lr.unsqueeze(0)
test_hr = test_hr.unsqueeze(0)
test_bc = test_bc.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 2 every 40 epochs
if (epoch + 1) % 40 == 0:
for param_group in self.optimizer.param_groups:
param_group['lr'] /= 2.0
print('Learning rate decay: lr={}'.format(
self.optimizer.param_groups[0]['lr']))
epoch_loss = 0
epoch_loss_D = 0
for iter, (lr, hr, bc_lr) in enumerate(train_data_loader):
# input data (low resolution image)
x_ = Variable(hr).cuda() if self.gpu_mode else Variable(hr)
y_ = Variable(lr).cuda() if self.gpu_mode else Variable(hr)
bc_y_ = Variable(bc_lr).cuda() if self.gpu_mode else Variable(
hr)
recon_image = self.model(y_)
recon_image = recon_image + bc_y_
loss_G = 0
loss_D = 0
style_loss = 0
loss_T = []
if 'A' in self.model_loss:
# ---------------------
# Train Discriminator
# ---------------------
self.valid = Variable(Tensor(
np.ones(x_.size()[0]).reshape((x_.size()[0], 1))),
requires_grad=False)
self.fake = Variable(Tensor(
np.zeros(x_.size()[0]).reshape((x_.size()[0], 1))),
requires_grad=False)
if self.gpu_mode:
self.valid.cuda()
self.fake.cuda()
self.optimizer_D.zero_grad()
# Loss of real and fake images
if self.loss_F == "BCEWithLogitsLoss":
loss_real = self.criterion_GAN(self.discriminator(x_),
self.valid)
loss_fake = self.criterion_GAN(
self.discriminator(recon_image.detach()),
self.fake)
elif self.loss_F == "MSE":
loss_real = self.mse_loss(self.discriminator(x_),
self.valid)
loss_fake = self.mse_loss(
self.discriminator(recon_image.detach()),
self.fake)
# Total loss
loss_D = (loss_real + loss_fake) / 2
loss_D.backward()
self.optimizer_D.step()
epoch_loss_D += loss_D.data[0]
if iter % self.D_period == 0:
# -----------------
# Train Generator
# -----------------
self.optimizer.zero_grad()
loss_a, loss_output_m2, loss_output_m5, style_score, loss_G, loss_T = Loss.loss_op(
self, self, recon_image, x_)
loss = (2*0.1*loss_output_m2) + \
(2*0.01*loss_output_m5) + \
loss_a*2 + style_score*1e-6
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
utils.print_loss(self, epoch, len(train_data_loader),
loss, style_score, loss_output_m2,
loss_output_m5, iter, loss_a, loss_D,
loss_G, loss_T)
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
del x_, y_, bc_y_
if self.gpu_mode:
torch.cuda.empty_cache()
else:
# update network
self.optimizer.zero_grad()
loss_a, loss_output_m2, loss_output_m5, style_score, loss_G, loss_T = Loss.loss_op(
self, self, recon_image, x_)
loss = (2*0.1*loss_output_m2) + \
(2*0.01*loss_output_m5) + \
loss_a*1e-3 + style_score*1e-6
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
#epoch_loss += loss
utils.print_loss(self, epoch, len(train_data_loader), loss,
style_score, loss_output_m2,
loss_output_m5, iter, loss_a, loss_D,
loss_G, loss_T)
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
del x_, y_, bc_y_
if self.gpu_mode:
torch.cuda.empty_cache()
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
y_ = Variable(test_lr).cuda() if self.gpu_mode else Variable(
test_lr)
bc = Variable(test_bc).cuda() if self.gpu_mode else Variable(
test_bc)
recon_img = self.model(y_)
recon_img = recon_img + bc
sr_img = recon_img[0].cpu().data
# save result image
save_dir = os.path.join(self.save_dir, 'train_result')
utils.save_img(sr_img,
epoch + 1,
save_dir=save_dir,
is_training=True)
print('Result image at epoch %d is saved.' % (epoch + 1))
#
utils.plot_loss([avg_loss], epoch, save_dir=save_dir)
if 'A' in self.model_loss:
avg_loss_D.append(epoch_loss_D / len(train_data_loader))
utils.plot_loss([avg_loss_D], epoch, save_dir=save_dir + '/D')
del y_, bc
if self.gpu_mode:
torch.cuda.empty_cache()
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
utils.save_model(self, epoch + 1)
test_v2.save_TrainingTest(self, epoch + 1)
# calculate psnrs
if self.num_channels == 1:
gt_img = test_hr[0][0].unsqueeze(0)
lr_img = test_lr[0][0].unsqueeze(0)
bc_img = test_bc[0][0].unsqueeze(0)
else:
gt_img = test_hr[0]
lr_img = test_lr[0]
bc_img = test_bc[0]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
self.num_epochs,
save_dir=save_dir,
is_training=True)
print('Training result image is saved.')
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=save_dir)
print('Training is finished.')
# Save final trained parameters of model
utils.save_model(self, epoch=None)
def test(self):
# networks
self.model = Net(3, 64, 3, 1)
# load model
utils.load_model(self, False, False)
# self.model.load_state_dict(torch.load('epoch0_70.pkl'))
if self.gpu_mode:
self.model.cuda()
# load dataset
for test_dataset in self.test_dataset:
test_data_loader = self.load_dataset(dataset=test_dataset,
is_train=False)
# Test
print('Test is started.')
img_num = 0
total_img_num = len(test_data_loader)
self.model.eval()
for lr, hr, bc, _ in test_data_loader:
# input data (low resolution image)
x_ = Variable(hr).cuda() if self.gpu_mode else Variable(hr)
y_ = Variable(lr).cuda() if self.gpu_mode else Variable(hr)
bc_ = Variable(bc).cuda() if self.gpu_mode else Variable(hr)
# prediction
recon_imgs = self.model(y_)
recon_imgs += bc_
for i, recon_img in enumerate(recon_imgs):
img_num += 1
sr_img = recon_img.cpu().data
# save result image
# save_dir = os.path.join(
# self.save_dir, 'test_result_texture', test_dataset)
save_dir = os.path.join(self.save_dir, 'test_result',
test_dataset)
utils.save_img(sr_img, img_num, save_dir=save_dir)
# utils.save_img(x_, img_num, save_dir=os.path.join(
# self.save_dir, 'test_HR', test_dataset))
# calculate psnrs
if self.num_channels == 1:
gt_img = hr[i][0].unsqueeze(0)
lr_img = lr[i][0].unsqueeze(0)
bc_img = bc[i][0].unsqueeze(0)
else:
gt_img = hr[i]
lr_img = lr[i]
bc_img = bc[i]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
img_num,
save_dir=save_dir)
print('Test DB: %s, Saving result images...[%d/%d]' %
(test_dataset, img_num, total_img_num))
print('Test is finishied.')
def train(self):
# networks
self.model = Net(num_channels=self.num_channels,
base_filter=64,
num_residuals=16)
# weigh initialization
self.model.weight_init()
# optimizer
self.optimizer = optim.Adam(self.model.parameters(),
lr=self.lr,
betas=(0.9, 0.999),
eps=1e-8)
# loss function
if self.gpu_mode:
self.model.cuda()
self.L1_loss = nn.L1Loss().cuda()
else:
self.L1_loss = nn.L1Loss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset=self.train_dataset,
is_train=True)
test_data_loader = self.load_dataset(dataset=self.test_dataset[0],
is_train=False)
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.makedirs(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_lr, test_hr, test_bc = test_data_loader.dataset.__getitem__(2)
test_lr = test_lr.unsqueeze(0)
test_hr = test_hr.unsqueeze(0)
test_bc = test_bc.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 2 every 40 epochs
if (epoch + 1) % 40 == 0:
for param_group in self.optimizer.param_groups:
param_group['lr'] /= 2.0
print('Learning rate decay: lr={}'.format(
self.optimizer.param_groups[0]['lr']))
epoch_loss = 0
for iter, (lr, hr, _) in enumerate(train_data_loader):
# input data (low resolution image)
if self.num_channels == 1:
x_ = Variable(hr[:, 0].unsqueeze(1))
y_ = Variable(lr[:, 0].unsqueeze(1))
else:
x_ = Variable(hr)
y_ = Variable(lr)
if self.gpu_mode:
x_ = x_.cuda()
y_ = y_.cuda()
# update network
self.optimizer.zero_grad()
recon_image = self.model(y_)
loss = self.L1_loss(recon_image, x_)
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
print('Epoch: [%2d] [%4d/%4d] loss: %.8f' %
((epoch + 1),
(iter + 1), len(train_data_loader), loss.data[0]))
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
if self.num_channels == 1:
y_ = Variable(test_lr[:, 0].unsqueeze(1))
else:
y_ = Variable(test_lr)
if self.gpu_mode:
y_ = y_.cuda()
recon_img = self.model(y_)
sr_img = recon_img[0].cpu().data
# save result image
save_dir = os.path.join(self.save_dir, 'train_result')
utils.save_img(sr_img,
epoch + 1,
save_dir=save_dir,
is_training=True)
print('Result image at epoch %d is saved.' % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# calculate psnrs
if self.num_channels == 1:
gt_img = test_hr[0][0].unsqueeze(0)
lr_img = test_lr[0][0].unsqueeze(0)
bc_img = test_bc[0][0].unsqueeze(0)
else:
gt_img = test_hr[0]
lr_img = test_lr[0]
bc_img = test_bc[0]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
self.num_epochs,
save_dir=save_dir,
is_training=True)
print('Training result image is saved.')
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=save_dir)
print('Training is finished.')
# Save final trained parameters of model
self.save_model(epoch=None)
def train(self):
# networks
self.model = Net(num_channels=self.num_channels, base_filter=64)
# weigh initialization
self.model.weight_init(mean=0.0, std=0.001)
# optimizer
self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr)
# loss function
if self.gpu_mode:
self.model.cuda()
self.MSE_loss = nn.MSELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.model)
print('----------------------------------------------')
# load dataset
train_data_loader = self.load_dataset(dataset='train')
test_data_loader = self.load_dataset(dataset='test')
# set the logger
log_dir = os.path.join(self.save_dir, 'logs')
if not os.path.exists(log_dir):
os.mkdir(log_dir)
logger = Logger(log_dir)
################# Train #################
print('Training is started.')
avg_loss = []
step = 0
# test image
test_input, test_target = test_data_loader.dataset.__getitem__(2)
test_input = test_input.unsqueeze(0)
test_target = test_target.unsqueeze(0)
self.model.train()
for epoch in range(self.num_epochs):
epoch_loss = 0
for iter, (input, target) in enumerate(train_data_loader):
# input data (bicubic interpolated image)
if self.gpu_mode:
# exclude border pixels from loss computation
x_ = Variable(utils.shave(target, border_size=8).cuda())
y_ = Variable(utils.img_interp(input, self.scale_factor).cuda())
else:
x_ = Variable(utils.shave(target, border_size=8))
y_ = Variable(utils.img_interp(input, self.scale_factor))
# update network
self.optimizer.zero_grad()
recon_image = self.model(y_)
loss = self.MSE_loss(recon_image, x_)
loss.backward()
self.optimizer.step()
# log
epoch_loss += loss.data[0]
print("Epoch: [%2d] [%4d/%4d] loss: %.8f" % ((epoch + 1), (iter + 1), len(train_data_loader), loss.data[0]))
# tensorboard logging
logger.scalar_summary('loss', loss.data[0], step + 1)
step += 1
# avg. loss per epoch
avg_loss.append(epoch_loss / len(train_data_loader))
# prediction
recon_imgs = self.model(Variable(utils.img_interp(test_input, self.scale_factor).cuda()))
recon_img = recon_imgs[0].cpu().data
gt_img = utils.shave(test_target[0], border_size=8)
lr_img = test_input[0]
bc_img = utils.shave(utils.img_interp(test_input[0], self.scale_factor), border_size=8)
# calculate psnrs
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(recon_img, gt_img)
# save result images
result_imgs = [gt_img, lr_img, bc_img, recon_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs, psnrs, epoch + 1, save_dir=self.save_dir, is_training=True)
print("Saving training result images at epoch %d" % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# Plot avg. loss
utils.plot_loss([avg_loss], self.num_epochs, save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
def train_ae(
net,
data_loader,
hyp,
criterion,
optimizer,
scheduler,
PATH="",
test_loader=None,
epoch_start=0,
epoch_end=1,
):
info_period = hyp["info_period"]
noiseSTD = hyp["noiseSTD"]
device = hyp["device"]
normalize = hyp["normalize"]
network = hyp["network"]
supervised = hyp["supervised"]
if normalize:
net.normalize()
if hyp["denoising"]:
if test_loader is not None:
with torch.no_grad():
psnr = []
for idx_test, (img_test, _) in tqdm(enumerate(test_loader)):
img_test_noisy = (img_test + noiseSTD / 255 *
torch.randn(img_test.shape)).to(device)
img_test_hat, _, _ = net(img_test_noisy)
img_test_noisy.detach()
psnr.append(
utils.PSNR(
img_test[0,
0, :, :].clone().detach().cpu().numpy(),
img_test_hat[
0, 0, :, :].clone().detach().cpu().numpy(),
))
np.save(os.path.join(PATH, "psnr_init.npy"), np.array(psnr))
print("PSNR: {}".format(np.round(np.array(psnr), decimals=4)))
for epoch in tqdm(range(epoch_start, epoch_end)):
scheduler.step()
loss_all = 0
for idx, (img, _) in tqdm(enumerate(data_loader)):
optimizer.zero_grad()
if supervised:
img = img.to(device)
noise = noiseSTD / 255 * torch.randn(img.shape).to(device)
noisy_img = img + noise
img_hat, _, _ = net(noisy_img)
loss = criterion(img, img_hat)
else:
noisy_img = (
img + noiseSTD / 255 * torch.randn(img.shape)).to(device)
img_hat, _, _ = net(noisy_img)
loss = criterion(noisy_img, img_hat)
loss_all += float(loss.item())
# ===================backward====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
if normalize:
net.normalize()
if idx % info_period == 0:
print("loss:{:.8f} ".format(loss.item()))
torch.cuda.empty_cache()
# ===================log========================
if hyp["denoising"]:
if test_loader is not None:
with torch.no_grad():
psnr = []
for idx_test, (img_test,
_) in tqdm(enumerate(test_loader)):
img_test_noisy = (
img_test + noiseSTD / 255 *
torch.randn(img_test.shape)).to(device)
img_test_hat, _, _ = net(img_test_noisy)
psnr.append(
utils.PSNR(
img_test[
0, 0, :, :].clone().detach().cpu().numpy(),
img_test_hat[
0, 0, :, :].clone().detach().cpu().numpy(),
))
np.save(
os.path.join(PATH, "psnr_epoch{}.npy".format(epoch)),
np.array(psnr),
)
print("")
print("PSNR: {}".format(
np.round(np.array(psnr), decimals=4)))
torch.save(loss_all, os.path.join(PATH,
"loss_epoch{}.pt".format(epoch)))
torch.save(net, os.path.join(PATH, "model_epoch{}.pt".format(epoch)))
print("epoch [{}/{}], loss:{:.8f} ".format(epoch + 1,
hyp["num_epochs"],
loss.item()))
return net
def train(self):
# load dataset
train_data_loader = self.load_dataset(dataset=self.train_dataset,
is_train=True)
test_data_loader = self.load_dataset(dataset=self.test_dataset[0],
is_train=False)
# networks
self.G = Generator(num_channels=self.num_channels,
base_filter=64,
num_residuals=16)
self.D = Discriminator(num_channels=self.num_channels,
base_filter=64,
image_size=self.crop_size)
# weigh initialization
self.G.weight_init()
self.D.weight_init()
# For the content loss
self.feature_extractor = FeatureExtractor(
models.vgg19(pretrained=True))
# optimizer
self.G_optimizer = optim.Adam(self.G.parameters(),
lr=self.lr,
betas=(0.9, 0.999))
# self.D_optimizer = optim.Adam(self.D.parameters(), lr=self.lr, betas=(0.9, 0.999))
self.D_optimizer = optim.SGD(self.D.parameters(),
lr=self.lr / 100,
momentum=0.9,
nesterov=True)
# loss function
if self.gpu_mode:
self.G.cuda()
self.D.cuda()
self.feature_extractor.cuda()
self.MSE_loss = nn.MSELoss().cuda()
self.BCE_loss = nn.BCELoss().cuda()
else:
self.MSE_loss = nn.MSELoss()
self.BCE_loss = nn.BCELoss()
print('---------- Networks architecture -------------')
utils.print_network(self.G)
utils.print_network(self.D)
print('----------------------------------------------')
# set the logger
G_log_dir = os.path.join(self.save_dir, 'G_logs')
if not os.path.exists(G_log_dir):
os.mkdir(G_log_dir)
G_logger = Logger(G_log_dir)
D_log_dir = os.path.join(self.save_dir, 'D_logs')
if not os.path.exists(D_log_dir):
os.mkdir(D_log_dir)
D_logger = Logger(D_log_dir)
################# Pre-train generator #################
self.epoch_pretrain = 50
# Load pre-trained parameters of generator
if not self.load_model(is_pretrain=True):
# Pre-training generator for 50 epochs
print('Pre-training is started.')
self.G.train()
for epoch in range(self.epoch_pretrain):
for iter, (lr, hr, _) in enumerate(train_data_loader):
# input data (low resolution image)
if self.num_channels == 1:
x_ = Variable(
utils.norm(hr[:, 0].unsqueeze(1), vgg=True))
y_ = Variable(
utils.norm(lr[:, 0].unsqueeze(1), vgg=True))
else:
x_ = Variable(utils.norm(hr, vgg=True))
y_ = Variable(utils.norm(lr, vgg=True))
if self.gpu_mode:
x_ = x_.cuda()
y_ = y_.cuda()
# Train generator
self.G_optimizer.zero_grad()
recon_image = self.G(y_)
# Content losses
content_loss = self.MSE_loss(recon_image, x_)
# Back propagation
G_loss_pretrain = content_loss
G_loss_pretrain.backward()
self.G_optimizer.step()
# log
print("Epoch: [%2d] [%4d/%4d] G_loss_pretrain: %.8f" %
((epoch + 1), (iter + 1), len(train_data_loader),
G_loss_pretrain.item()))
print('Pre-training is finished.')
# Save pre-trained parameters of generator
self.save_model(is_pretrain=True)
################# Adversarial train #################
print('Training is started.')
# Avg. losses
G_avg_loss = []
D_avg_loss = []
step = 0
# test image
test_lr, test_hr, test_bc = test_data_loader.dataset.__getitem__(2)
test_lr = test_lr.unsqueeze(0)
test_hr = test_hr.unsqueeze(0)
test_bc = test_bc.unsqueeze(0)
self.G.train()
self.D.train()
for epoch in range(self.num_epochs):
# learning rate is decayed by a factor of 10 every 20 epoch
if (epoch + 1) % 20 == 0:
for param_group in self.G_optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay for G: lr={}".format(
self.G_optimizer.param_groups[0]["lr"]))
for param_group in self.D_optimizer.param_groups:
param_group["lr"] /= 10.0
print("Learning rate decay for D: lr={}".format(
self.D_optimizer.param_groups[0]["lr"]))
G_epoch_loss = 0
D_epoch_loss = 0
for iter, (lr, hr, _) in enumerate(train_data_loader):
# input data (low resolution image)
mini_batch = lr.size()[0]
if self.num_channels == 1:
x_ = Variable(utils.norm(hr[:, 0].unsqueeze(1), vgg=True))
y_ = Variable(utils.norm(lr[:, 0].unsqueeze(1), vgg=True))
else:
x_ = Variable(utils.norm(hr, vgg=True))
y_ = Variable(utils.norm(lr, vgg=True))
if self.gpu_mode:
x_ = x_.cuda()
y_ = y_.cuda()
# labels
real_label = Variable(torch.ones(mini_batch).cuda())
fake_label = Variable(torch.zeros(mini_batch).cuda())
else:
# labels
real_label = Variable(torch.ones(mini_batch))
fake_label = Variable(torch.zeros(mini_batch))
# Reset gradient
self.D_optimizer.zero_grad()
# Train discriminator with real data
D_real_decision = self.D(x_)
D_real_loss = self.BCE_loss(D_real_decision[:, 0], real_label)
# Train discriminator with fake data
recon_image = self.G(y_)
D_fake_decision = self.D(recon_image)
D_fake_loss = self.BCE_loss(D_fake_decision[:, 0], fake_label)
D_loss = D_real_loss + D_fake_loss
# Back propagation
D_loss.backward()
self.D_optimizer.step()
# Reset gradient
self.G_optimizer.zero_grad()
# Train generator
recon_image = self.G(y_)
D_fake_decision = self.D(recon_image)
# Adversarial loss
GAN_loss = self.BCE_loss(D_fake_decision[:, 0], real_label)
# Content losses
mse_loss = self.MSE_loss(recon_image, x_)
x_VGG = Variable(utils.norm(hr, vgg=True).cuda())
recon_VGG = Variable(
utils.norm(recon_image.data, vgg=True).cuda())
real_feature = self.feature_extractor(x_VGG)
fake_feature = self.feature_extractor(recon_VGG)
vgg_loss = self.MSE_loss(fake_feature, real_feature.detach())
# Back propagation
G_loss = mse_loss + 6e-3 * vgg_loss + 1e-3 * GAN_loss
G_loss.backward()
self.G_optimizer.step()
# log
G_epoch_loss += G_loss.item()
D_epoch_loss += D_loss.item()
print("Epoch: [%2d] [%4d/%4d] G_loss: %.8f, D_loss: %.8f" %
((epoch + 1), (iter + 1), len(train_data_loader),
G_loss.item(), D_loss.item()))
# tensorboard logging
#G_logger.scalar_summary('losses', G_loss.item(), step + 1)
#D_logger.scalar_summary('losses', D_loss.item(), step + 1)
step += 1
# avg. loss per epoch
G_avg_loss.append(G_epoch_loss / len(train_data_loader))
D_avg_loss.append(D_epoch_loss / len(train_data_loader))
# prediction
if self.num_channels == 1:
y_ = Variable(utils.norm(test_lr[:, 0].unsqueeze(1), vgg=True))
else:
y_ = Variable(utils.norm(test_lr, vgg=True))
if self.gpu_mode:
y_ = y_.cuda()
recon_img = self.G(y_)
sr_img = utils.denorm(recon_img[0].cpu().data, vgg=True)
# save result image
save_dir = os.path.join(self.save_dir, 'train_result')
utils.save_img(sr_img,
epoch + 1,
save_dir=save_dir,
is_training=True)
print('Result image at epoch %d is saved.' % (epoch + 1))
# Save trained parameters of model
if (epoch + 1) % self.save_epochs == 0:
self.save_model(epoch + 1)
# calculate psnrs
if self.num_channels == 1:
gt_img = test_hr[0][0].unsqueeze(0)
lr_img = test_lr[0][0].unsqueeze(0)
bc_img = test_bc[0][0].unsqueeze(0)
else:
gt_img = test_hr[0]
lr_img = test_lr[0]
bc_img = test_bc[0]
bc_psnr = utils.PSNR(bc_img, gt_img)
recon_psnr = utils.PSNR(sr_img, gt_img)
# plot result images
result_imgs = [gt_img, lr_img, bc_img, sr_img]
psnrs = [None, None, bc_psnr, recon_psnr]
utils.plot_test_result(result_imgs,
psnrs,
self.num_epochs,
save_dir=save_dir,
is_training=True)
print('Training result image is saved.')
# Plot avg. loss
utils.plot_loss([G_avg_loss, D_avg_loss],
self.num_epochs,
save_dir=self.save_dir)
print("Training is finished.")
# Save final trained parameters of model
self.save_model(epoch=None)
def main():
# ==========
# parameters
# ==========
opts_dict = receive_arg()
unit = opts_dict['test']['criterion']['unit']
# ==========
# open logger
# ==========
log_fp = open(opts_dict['train']['log_path'], 'w')
msg = (
f"{'<' * 10} Test {'>' * 10}\n"
f"Timestamp: [{utils.get_timestr()}]\n"
f"\n{'<' * 10} Options {'>' * 10}\n"
f"{utils.dict2str(opts_dict['test'])}"
)
print(msg)
log_fp.write(msg + '\n')
log_fp.flush()
# ==========
# Ensure reproducibility or Speed up
# ==========
#torch.backends.cudnn.benchmark = False # if reproduce
#torch.backends.cudnn.deterministic = True # if reproduce
torch.backends.cudnn.benchmark = True # speed up
# ==========
# create test data prefetchers
# ==========
# create datasets
test_ds_type = opts_dict['dataset']['test']['type']
radius = opts_dict['network']['radius']
assert test_ds_type in dataset.__all__, \
"Not implemented!"
test_ds_cls = getattr(dataset, test_ds_type)
test_ds = test_ds_cls(
opts_dict=opts_dict['dataset']['test'],
radius=radius
)
test_num = len(test_ds)
test_vid_num = test_ds.get_vid_num()
# create datasamplers
test_sampler = None # no need to sample test data
# create dataloaders
test_loader = utils.create_dataloader(
dataset=test_ds,
opts_dict=opts_dict,
sampler=test_sampler,
phase='val'
)
assert test_loader is not None
# create dataloader prefetchers
test_prefetcher = utils.CPUPrefetcher(test_loader)
# ==========
# create & load model
# ==========
model = MFVQE(opts_dict=opts_dict['network'])
checkpoint_save_path = opts_dict['test']['checkpoint_save_path']
msg = f'loading model {checkpoint_save_path}...'
print(msg)
log_fp.write(msg + '\n')
checkpoint = torch.load(checkpoint_save_path)
if 'module.' in list(checkpoint['state_dict'].keys())[0]: # multi-gpu training
new_state_dict = OrderedDict()
for k, v in checkpoint['state_dict'].items():
name = k[7:] # remove module
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
else: # single-gpu training
model.load_state_dict(checkpoint['state_dict'])
msg = f'> model {checkpoint_save_path} loaded.'
print(msg)
log_fp.write(msg + '\n')
model = model.cuda()
model.eval()
# ==========
# define criterion
# ==========
# define criterion
assert opts_dict['test']['criterion'].pop('type') == \
'PSNR', "Not implemented."
criterion = utils.PSNR()
# ==========
# validation
# ==========
# create timer
total_timer = utils.Timer()
# create counters
per_aver_dict = dict()
ori_aver_dict = dict()
name_vid_dict = dict()
for index_vid in range(test_vid_num):
per_aver_dict[index_vid] = utils.Counter()
ori_aver_dict[index_vid] = utils.Counter()
name_vid_dict[index_vid] = ""
pbar = tqdm(
total=test_num,
ncols=opts_dict['test']['pbar_len']
)
# fetch the first batch
test_prefetcher.reset()
val_data = test_prefetcher.next()
with torch.no_grad():
while val_data is not None:
# get data
gt_data = val_data['gt'].cuda() # (B [RGB] H W)
lq_data = val_data['lq'].cuda() # (B T [RGB] H W)
index_vid = val_data['index_vid'].item()
name_vid = val_data['name_vid'][0] # bs must be 1!
b, _, c, _, _ = lq_data.shape
assert b == 1, "Not supported!"
input_data = torch.cat(
[lq_data[:,:,i,...] for i in range(c)],
dim=1
) # B [R1 ... R7 G1 ... G7 B1 ... B7] H W
enhanced_data = model(input_data) # (B [RGB] H W)
# eval
batch_ori = criterion(lq_data[0, radius, ...], gt_data[0])
batch_perf = criterion(enhanced_data[0], gt_data[0])
# display
pbar.set_description(
"{:s}: [{:.3f}] {:s} -> [{:.3f}] {:s}"
.format(name_vid, batch_ori, unit, batch_perf, unit)
)
pbar.update()
# log
per_aver_dict[index_vid].accum(volume=batch_perf)
ori_aver_dict[index_vid].accum(volume=batch_ori)
if name_vid_dict[index_vid] == "":
name_vid_dict[index_vid] = name_vid
else:
assert name_vid_dict[index_vid] == name_vid, "Something wrong."
# fetch next batch
val_data = test_prefetcher.next()
# end of val
pbar.close()
# log
msg = '\n' + '<' * 10 + ' Results ' + '>' * 10
print(msg)
log_fp.write(msg + '\n')
for index_vid in range(test_vid_num):
per = per_aver_dict[index_vid].get_ave()
ori = ori_aver_dict[index_vid].get_ave()
name_vid = name_vid_dict[index_vid]
msg = "{:s}: [{:.3f}] {:s} -> [{:.3f}] {:s}".format(
name_vid, ori, unit, per, unit
)
print(msg)
log_fp.write(msg + '\n')
ave_per = np.mean([
per_aver_dict[index_vid].get_ave() for index_vid in range(test_vid_num)
])
ave_ori = np.mean([
ori_aver_dict[index_vid].get_ave() for index_vid in range(test_vid_num)
])
msg = (
f"{'> ori: [{:.3f}] {:s}'.format(ave_ori, unit)}\n"
f"{'> ave: [{:.3f}] {:s}'.format(ave_per, unit)}\n"
f"{'> delta: [{:.3f}] {:s}'.format(ave_per - ave_ori, unit)}"
)
print(msg)
log_fp.write(msg + '\n')
log_fp.flush()
# ==========
# final log & close logger
# ==========
total_time = total_timer.get_interval() / 3600
msg = "TOTAL TIME: [{:.1f}] h".format(total_time)
print(msg)
log_fp.write(msg + '\n')
msg = (
f"\n{'<' * 10} Goodbye {'>' * 10}\n"
f"Timestamp: [{utils.get_timestr()}]"
)
print(msg)
log_fp.write(msg + '\n')
log_fp.close()
def train_ae(
net,
data_loader,
hyp,
criterion,
optimizer,
scheduler,
PATH="",
test_loader=None,
epoch_start=0,
epoch_end=1,
H_true=None,
org_img=None,
):
info_period = hyp["info_period"]
device = hyp["device"]
zero_mean_filters = hyp["zero_mean_filters"]
normalize = hyp["normalize"]
network = hyp["network"]
mu = hyp["mu"]
supervised = hyp["supervised"]
Jg = hyp["num_trials"]
data_distribution = hyp["data_distribution"]
model_distribution = hyp["model_distribution"]
peak = hyp["peak"]
train_L = hyp["train_L"]
if data_distribution == "gaussian":
noiseSTD = hyp["noiseSTD"]
if normalize:
net.normalize()
if hyp["denoising"]:
if test_loader is not None:
with torch.no_grad():
psnr = []
t = 0
for idx_test, (img_test, _) in enumerate(test_loader):
t += 1
psnr_i = 0
N = 20
for _ in range(N):
if data_distribution == "gaussian":
img_test_noisy = (
img_test + noiseSTD / 255 * torch.randn(img_test.shape)
).to(device)
elif data_distribution == "binomial":
img_test = img_test.to(device)
sampler = torch.distributions.bernoulli.Bernoulli(
probs=img_test
)
img_test_noisy = sampler.sample()
for j in range(Jg - 1):
img_test_noisy += sampler.sample()
img_test_noisy /= Jg
elif data_distribution == "poisson":
if img_test[img_test > 0] is not None:
img_test[img_test == 0] = torch.min(
img_test[img_test > 0]
)
img_test = img_test.to(device)
Q = torch.max(img_test) / peak
rate = img_test / Q
if torch.isnan(torch.min(rate)):
continue
sampler = torch.distributions.poisson.Poisson(rate)
if model_distribution == "poisson":
img_test_noisy = sampler.sample()
else:
img_test_noisy = sampler.sample() * Q
Hx_hat, _, _ = net(img_test_noisy, mu)
if model_distribution == "gaussian":
img_test_hat = Hx_hat + mu
elif model_distribution == "binomial":
img_test_hat = torch.nn.Sigmoid()(Hx_hat + mu)
elif model_distribution == "poisson":
img_test_hat = Q * torch.exp(Hx_hat + mu)
psnr_i += utils.PSNR(
img_test[0, 0, :, :].detach().cpu().numpy(),
img_test_hat[0, 0, :, :].detach().cpu().numpy(),
)
psnr.append(psnr_i / N)
if model_distribution == "poisson":
noisy_psnr = utils.PSNR(
img_test[0, 0, :, :].detach().cpu().numpy(),
(Q * img_test_noisy[0, 0, :, :]).detach().cpu().numpy(),
)
else:
noisy_psnr = utils.PSNR(
img_test[0, 0, :, :].detach().cpu().numpy(),
img_test_noisy[0, 0, :, :].detach().cpu().numpy(),
)
np.save(os.path.join(PATH, "psnr_init.npy"), np.array(psnr))
print(
"PSNR: {}, {}".format(
np.round(np.array(noisy_psnr), decimals=4),
np.round(np.array(psnr), decimals=4),
)
)
nan_ctr = 0
for epoch in tqdm(range(epoch_start, epoch_end)):
scheduler.step()
loss_all = 0
for idx, (img, _) in tqdm(enumerate(data_loader)):
optimizer.zero_grad()
if data_distribution == "gaussian":
img_noisy = (img + noiseSTD / 255 * torch.randn(img.shape)).to(device)
img = img.to(device)
elif data_distribution == "binomial":
img = img.to(device)
sampler = torch.distributions.bernoulli.Bernoulli(probs=img)
img_noisy = sampler.sample()
for j in range(Jg - 1):
img_noisy += sampler.sample()
img_noisy /= Jg
elif data_distribution == "poisson":
if torch.sum(img) != 0:
img[img == 0] = torch.min(img[img > 0])
else:
continue
img = img.to(device)
Q = torch.max(img) / peak
rate = img / Q
if torch.isnan(torch.min(rate)):
continue
sampler = torch.distributions.poisson.Poisson(rate)
if model_distribution == "poisson":
img_noisy = sampler.sample()
else:
img_noisy = sampler.sample() * Q
if torch.isnan(torch.sum(img_noisy)):
print("img_noisy got NaN!")
continue
Hx, x_hat, _ = net(img_noisy, mu)
if torch.isnan(torch.sum(Hx)):
print("Hx got NaN!")
nan_ctr += 1
if nan_ctr > 20:
break
else:
continue
if supervised:
if model_distribution == "poisson":
loss = criterion(img, Hx, Q)
else:
loss = criterion(img, Hx)
else:
loss = criterion(img_noisy, Hx)
if loss > 10:
f.write("skip. loss is large! \r\n")
print("skip. loss is large!")
continue
loss_all += float(loss.item())
# ===================backward====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
if zero_mean_filters:
net.zero_mean()
if normalize:
net.normalize()
if idx % info_period == 0:
if H_true is not None:
err_H = utils.err2d_H(H_true, net.get_param("H"))
print(
"loss:{:.4f}, err_H:{:4f}\n".format(loss.item(), np.mean(err_H))
)
else:
print("loss:{:.4f}\n".format(loss.item()))
torch.cuda.empty_cache()
# ===================log========================
if hyp["denoising"]:
if test_loader is not None:
with torch.no_grad():
psnr = []
t = 0
for idx_test, (img_test, _) in enumerate(test_loader):
t += 1
psnr_i = 0
N = 20
for _ in range(N):
if data_distribution == "gaussian":
img_test_noisy = (
img_test
+ noiseSTD / 255 * torch.randn(img_test.shape)
).to(device)
elif data_distribution == "binomial":
img_test = img_test.to(device)
sampler = torch.distributions.bernoulli.Bernoulli(
probs=img_test
)
img_test_noisy = sampler.sample()
for j in range(Jg - 1):
img_test_noisy += sampler.sample()
img_test_noisy /= Jg
elif data_distribution == "poisson":
if torch.sum(img_test) != 0:
img_test[img_test == 0] = torch.min(
img_test[img_test > 0]
)
img_test = img_test.to(device)
Q = torch.max(img_test) / peak
rate = img_test / Q
if torch.isnan(torch.min(rate)):
continue
sampler = torch.distributions.poisson.Poisson(rate)
if model_distribution == "poisson":
img_test_noisy = sampler.sample()
else:
img_test_noisy = sampler.sample() * Q
Hx_hat, _, _ = net(img_test_noisy, mu)
if model_distribution == "gaussian":
img_test_hat = Hx_hat + mu
elif model_distribution == "binomial":
img_test_hat = torch.nn.Sigmoid()(Hx_hat + mu)
elif model_distribution == "poisson":
img_test_hat = Q * torch.exp(Hx_hat + mu)
psnr_i += utils.PSNR(
img_test[0, 0, :, :].detach().cpu().numpy(),
img_test_hat[0, 0, :, :].detach().cpu().numpy(),
)
psnr.append(psnr_i / N)
np.save(
os.path.join(PATH, "psnr_epoch{}.npy".format(epoch)),
np.array(psnr),
)
print("PSNR: {}".format(np.round(np.array(psnr), decimals=4)))
torch.save(loss_all, os.path.join(PATH, "loss_epoch{}.pt".format(epoch)))
torch.save(net, os.path.join(PATH, "model_epoch{}.pt".format(epoch)))
print(
"epoch [{}/{}], loss:{:.4f} ".format(
epoch + 1, hyp["num_epochs"], loss.item()
)
)
if torch.isnan(torch.min(net.get_param("H"))):
print("network got NaN!")
break
return net
#utils.read_imagecollection('/home/wdd/pytorch-vdsr-3d/3D_VDSR_001')
for count_d in range((num // batch_generate_size)):
for count_h in range((h // batch_generate_size)):
for count_w in range((w // batch_generate_size)):
pixel_start_d = count_d * batch_generate_size
pixel_end_d = (count_d + 1) * batch_generate_size
pixel_start_h = count_h * batch_generate_size
pixel_end_h = (count_h + 1) * batch_generate_size
pixel_start_w = count_w * batch_generate_size
pixel_end_w = (count_w + 1) * batch_generate_size
sub_block = utils.read_imagecollection(
file_path=default_path + str(count_d) + str(count_h) +
str(count_w),
image_format='bmp')
# os.remove('3D_VDSR_'+str(count_d)+str(count_h)+str(count_w))
#print ('current path is:'+'3D_VDSR_'+str(count_d)+str(count_h)+str(count_w),'shape of sub_block is:',sub_block.shape)
reconstruction_output[pixel_start_d:pixel_end_d,
pixel_start_h:pixel_end_h,
pixel_start_w:pixel_end_w] = sub_block
dataset_ori = utils.read_imagecollection(opt.oripath)
dataset_interp = utils.read_imagecollection(
opt.interpath, image_format='bmp') #读取interp path 路径,文件后缀要注意bmp或是jpg！！
#print ('======>Read original image from : ',opt.oripath,' Read low resolution images from : ',opt.interpath)
#print ('PSNR of interp:',utils.PSNR(dataset_interp[:reconstruction_size,:reconstruction_size,:reconstruction_size],dataset_ori[:reconstruction_size,:reconstruction_size,:reconstruction_size]))
# print ('PSNR of reconstructor:',utils.PSNR(reconstruction_output,dataset_ori[:reconstruction_size,:reconstruction_size,:reconstruction_size]))
print(utils.PSNR(reconstruction_output, dataset_ori[:num, :h, :w]))
# Image.fromarray(reconstruction_output[..., 100]).show()
# Image.fromarray(dataset_interp[..., 100]).show()
utils.generate_2Dimage(save_mode=os.path.join(default_path, sub_name),
array_like=reconstruction_output)
