def compare_pois_l2(s, c, saved_epoch, img_ind, add2exp='_400/',\
path2valdata='./DATASETS/BSDS500/BSDS500_validation_MAXVALs_01_2/'):
exp_name_pois = 'pois' + add2exp + 's{}c{}'.format(s, c)
exp_name_l2 = 'l2' + add2exp + 's{}c{}'.format(s, c)
model_pois = PoisNet(output_features=c, stages=s)
model_l2 = UDNet(output_features=c, stages=s)
path2dataset = path2valdata
BSDSval = BSDS500(path2dataset + 'val/', get_name=True)
gt, noisy, file_name = BSDSval[img_ind]
split = file_name.split('_')
name, maxval = split[0], split[2][7:]
maxval = (int(maxval[0]) * 10 + int(maxval[-1])) / 10
gt.unsqueeze_(0)
noisy.unsqueeze_(0)
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name_pois+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model_pois.load_state_dict(state['model_state_dict'])
estim_pois = model_pois(noisy, noisy).detach()
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name_l2+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model_l2.load_state_dict(state['model_state_dict'])
stdn = th.Tensor([5])
estim_l2 = model_l2(noisy, stdn, noisy).detach()
psnr_noisy = psnr(gt, noisy)
psnr_est_pois = psnr(gt, estim_pois)
psnr_est_l2 = psnr(gt, estim_l2)
gt_title = 'clear ({} in BSDS val)'.format(img_ind)
noisy_title = 'noisy (max val={}) \nPSNR: {:.2f} dB'.format(
maxval, psnr_noisy)
estim_pois_title = '{} (epoch={}),\nPSNR: {:.3f} dB'\
.format('pois', saved_epoch, psnr_est_pois)
estim_l2_title = '{} (epoch={}),\nPSNR: {:.3f} dB'\
.format('l2', saved_epoch, psnr_est_l2)
show_images([noisy, gt, estim_pois, estim_l2], \
[noisy_title, gt_title, estim_pois_title, estim_l2_title])
def do_inference_return_lists(noisy, gt, s, c, exp_name, saved_epoch, model_type='pois',\
app=None, prox_param=False, sharing_weights=True):
if model_type == 'pois':
model = PoisNet(output_features=c,
stages=s,
prox_param=prox_param,
convWeightSharing=sharing_weights)
elif model_type == 'l2':
model = UDNet(output_features=c,
stages=s,
alpha=prox_param,
convWeightSharing=sharing_weights)
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model.load_state_dict(state['model_state_dict'])
if model_type == 'pois':
estim = model(noisy, noisy).detach()
elif model_type == 'l2':
stdn = th.Tensor([1])
estim = model(noisy, stdn, noisy).detach()
psnr_est = psnr(gt, estim)
if app is not None:
estim_title = '{} (epoch={}),\nPSNR: {:.3f} dB'.format(
app, saved_epoch, psnr_est)
else:
estim_title = 'estim (epoch={}),\nPSNR: {:.3f} dB'.format(
saved_epoch, psnr_est)
return estim, estim_title
output_features = opt.channels
prox_param = opt.prox_param
sharing_weights = not opt.no_sharing_weights
print(prox_param)
print(sharing_weights)
model_type = opt.model_type
print(model_type)
if model_type == 'pois':
model = PoisNet(stages=stages, output_features=output_features,\
prox_param=prox_param, convWeightSharing=sharing_weights).cuda()
elif model_type == 'l2':
model = UDNet(stages=stages, output_features=output_features, \
alpha=prox_param, convWeightSharing=sharing_weights).cuda()
lr0 = opt.lr0
optimizer = Adam(model.parameters(), lr=lr0)
milestones = list(map(int, opt.milestones.strip('[]').split(','))) \
if opt.milestones != '' else []
scheduler = MultiStepLR(optimizer, milestones=milestones, \
gamma=0.1)
loss = opt.loss
if loss == 'MSE':
criterion = MSELoss().cuda()
elif loss == 'pois':
criterion = poisLLHLoss
elif loss == 'L1':
train_set.stdn)
print('===> Building model')
# Parameters that we need to specify in order to initialize our model
params = OrderedDict(kernel_size=opt.kernel_size,input_channels=input_channels,\
output_features=output_features,rbf_mixtures=opt.rbf_mixtures,\
rbf_precision=opt.rbf_precision,stages=opt.stages,pad=opt.pad,\
padType=opt.padType,convWeightSharing=opt.convWeightSharing,\
scale_f=opt.scale_f,scale_t=opt.scale_t,normalizedWeights=\
opt.normalizedWeights,zeroMeanWeights=opt.zeroMeanWeights,rbf_start=\
opt.rbf_start,rbf_end=opt.rbf_end,data_min=opt.data_min,data_max=\
opt.data_max,data_step=opt.data_step,alpha=opt.alpha,clb=opt.clb,\
cub=opt.cub)
model = UDNet(*params.values())
if opt.initModelPath != '':
state = th.load(opt.initModelPath,
map_location=lambda storage, loc: storage)
model.load_state_dict(state['model_state_dict'])
opt.resume = False
#criterion = nn.MSELoss(size_average=True,reduce=True)
criterion = PSNRLoss(peakval=opt.cub)
optimizer = optim.Adam(model.parameters(),
lr=opt.lr,
betas=(0.9, 0.999),
eps=1e-04)
def compare_pois_l2_pois_w_prox(s, c, saved_epoch, img_ind, add2exp='_400/',\
path2valdata='./DATASETS/BSDS500/BSDS500_validation_MAXVALs_01_2/'):
exp_name_pois = 'pois' + add2exp + 's{}c{}'.format(s, c)
exp_name_poisprox = 'pois_w_prox' + add2exp + 's{}c{}'.format(s, c)
exp_name_l2 = 'l2' + add2exp + 's{}c{}'.format(s, c)
model_pois = PoisNet(output_features=c, stages=s)
model_poisprox = PoisNet(output_features=c, stages=s, prox_param=True)
model_l2 = UDNet(output_features=c, stages=s)
path2dataset = path2valdata
BSDSval = BSDS500(path2dataset + 'val/', get_name=True)
gt, noisy, file_name = BSDSval[img_ind]
split = file_name.split('_')
name, maxval = split[0], split[2][7:]
maxval = (int(maxval[0]) * 10 + int(maxval[-1])) / 10
gt.unsqueeze_(0)
noisy.unsqueeze_(0)
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name_pois+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model_pois.load_state_dict(state['model_state_dict'])
estim_pois = model_pois(noisy, noisy).detach()
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name_poisprox+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model_poisprox.load_state_dict(state['model_state_dict'])
estim_poisprox = model_poisprox(noisy, noisy).detach()
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name_l2+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model_l2.load_state_dict(state['model_state_dict'])
stdn = th.Tensor([5])
estim_l2 = model_l2(noisy, stdn, noisy).detach()
psnr_noisy = psnr(gt, noisy)
psnr_est_pois = psnr(gt, estim_pois)
psnr_est_poisprox = psnr(gt, estim_poisprox)
psnr_est_l2 = psnr(gt, estim_l2)
gt_title = 'clear ({} in BSDS val)'.format(img_ind)
noisy_title = 'noisy (max val={}) \nPSNR: {:.2f} dB'.format(
maxval, psnr_noisy)
estim_pois_title = '{} (epoch={}),\nPSNR: {:.3f} dB'\
.format('pois', saved_epoch, psnr_est_pois)
estim_poisprox_title = '{} (epoch={}),\nPSNR: {:.3f} dB'\
.format('poisprox', saved_epoch, psnr_est_poisprox)
estim_l2_title = '{} (epoch={}),\nPSNR: {:.3f} dB'\
.format('l2', saved_epoch, psnr_est_l2)
images = [noisy, gt, estim_pois, estim_poisprox, estim_l2]
titles = [
noisy_title, gt_title, estim_pois_title, estim_poisprox_title,
estim_l2_title
]
fontsize = 15
images_corrected_dims = []
for i, img in enumerate(images):
if img.dim() == 4:
img = img[0]
img = img[0] if img.size()[0] == 1 else img.permute(1, 2, 0)
images_corrected_dims.append(img)
images = images_corrected_dims
figsize = (20, 10)
fig, ax = plt.subplots(2, 3, figsize=figsize)
fig.patch.set_facecolor('white')
ax[0, 0].imshow(images[0], cmap='gray')
ax[0, 0].set_axis_off()
ax[0, 0].set_title(noisy_title, fontsize=fontsize)
ax[0, 1].imshow(images[1], cmap='gray')
ax[0, 1].set_axis_off()
ax[0, 1].set_title(gt_title, fontsize=fontsize)
ax[1, 0].imshow(images[2], cmap='gray')
ax[1, 0].set_axis_off()
ax[1, 0].set_title(estim_pois_title, fontsize=fontsize)
ax[1, 1].imshow(images[3], cmap='gray')
ax[1, 1].set_axis_off()
ax[1, 1].set_title(estim_poisprox_title, fontsize=fontsize)
ax[1, 2].imshow(images[4], cmap='gray')
ax[1, 2].set_axis_off()
ax[1, 2].set_title(estim_l2_title, fontsize=fontsize)
ax[0, 2].remove()
fig.tight_layout()
def do_inference(s, c, exp_name, saved_epoch, model_type='pois', \
img_ind=None, img_from_train_dataset=False, clear_ind=0, app=None,\
path2valdata='./DATASETS/BSDS500/BSDS500_validation_MAXVALs_01_2/', prox_param=False, sharing_weights=True, do_VST=False):
if model_type == 'pois':
model = PoisNet(output_features=c,
stages=s,
prox_param=prox_param,
convWeightSharing=sharing_weights)
elif model_type == 'l2':
model = UDNet(output_features=c,
stages=s,
alpha=prox_param,
convWeightSharing=sharing_weights)
path2dataset = path2valdata
BSDSval = BSDS500(path2dataset+'val/', get_name=True, do_VST_4_visual=do_VST) \
if not img_from_train_dataset \
else BSDS500(path2dataset+'train/', get_name=True, do_VST_4_visual=do_VST)
if img_ind is None:
img_ind = 80
if do_VST:
gt, noisy, noisy_initial, file_name = BSDSval[img_ind]
else:
gt, noisy, file_name = BSDSval[img_ind]
split = file_name.split('_')
name, maxval = split[0], split[2][7:]
maxval = (int(maxval[0]) * 10 + int(maxval[-1])) / 10
gt.unsqueeze_(0)
noisy.unsqueeze_(0)
state = th.load('./PoisDenoiser/networks/PoisNet/models/'\
+exp_name+'/state_{}.pth'.format(saved_epoch),\
map_location=lambda storage,loc:storage)
model.load_state_dict(state['model_state_dict'])
if model_type == 'pois':
estim = model(noisy, noisy).detach()
elif model_type == 'l2':
stdn = th.Tensor([1])
estim = model(noisy, stdn, noisy).detach()
psnr_noisy = psnr(gt, noisy_initial) if do_VST else psnr(gt, noisy)
psnr_est = psnr(
gt, VST_backward_unbiased_exact(estim)) if do_VST else psnr(gt, estim)
gt_title = 'clear ({} in BSDS val)'.format(img_ind)
noisy_title = 'noisy (max val={}) \nPSNR: {:.2f} dB'.format(
maxval, psnr_noisy)
if app is not None:
estim_title = '{} (epoch={}),\nPSNR: {:.3f} dB'.format(
app, saved_epoch, psnr_est)
else:
estim_title = 'estim (epoch={}),\nPSNR: {:.3f} dB'.format(
saved_epoch, psnr_est)
if clear_ind == 0:
show_images([gt, noisy, estim], [gt_title, noisy_title, estim_title])
else:
show_images([noisy, gt, estim], [noisy_title, gt_title, estim_title])
print('===> Building model')
Lmodel = nn.ModuleList()
params = OrderedDict(kernel_size=opt.kernel_size,input_channels=input_channels,\
output_features=output_features,rbf_mixtures=opt.rbf_mixtures,\
rbf_precision=opt.rbf_precision,stages=1,pad=opt.pad,\
padType=opt.padType,convWeightSharing=opt.convWeightSharing,\
scale_f=opt.scale_f,scale_t=opt.scale_t,normalizedWeights=\
opt.normalizedWeights,zeroMeanWeights=opt.zeroMeanWeights,rbf_start=\
opt.rbf_start,rbf_end=opt.rbf_end,data_min=opt.data_min,data_max=\
opt.data_max,data_step=opt.data_step,alpha=opt.alpha,clb=opt.clb,\
cub=opt.cub)
for i in range(opt.stages):
Lmodel.append(UDNet(*params.values()))
tic()
for stage in range(opt.stages):
training_data_loader = DataLoader(dataset=train_set,num_workers=opt.threads,\
batch_size=opt.batchSize,shuffle=True)
testing_data_loader = DataLoader(dataset=test_set,num_workers=opt.threads,\
batch_size=opt.testBatchSize,shuffle=False)
#criterion = nn.MSELoss(size_average=True,reduce=True)
criterion = PSNRLoss(peakval=opt.cub)
stagePath = os.path.join(dirPath, 'stage' + str(stage + 1))
os.makedirs(stagePath, exist_ok=True)
smodel = Lmodel[stage]
optimizer = optim.Adam(smodel.parameters(),
lr=opt.lr,