def train_epoch(epoch):
epoch_loss_g = 0
total_psnr = 0
for i, batch in enumerate(training_data_loader, 1):
if i > max_train_iter:
break
target = Variable(batch[0])
input = Variable(batch[1])
if cuda:
target = target.cuda()
input = input.cuda()
# Train Generator
optimizerG.zero_grad()
reconstructed = gennet(input)
content_loss = content_criterion(reconstructed, target)
loss = content_loss
loss.backward()
optimizerG.step()
epoch_loss_g += loss.data[0]
total_psnr += psnr(loss.data[0])
print("Epoch[{}]({}/{}): Loss(G): {:.4f}".format(
epoch, i, max_train_iter, loss.data[0]))
avg_loss_g = epoch_loss_g / max_train_iter
avg_psnr = total_psnr / max_train_iter
print("Epoch {} : Avg Loss(G): {:.4f}".format(epoch, avg_loss_g))
plots.update([epoch, avg_psnr], rind=1, cind=1, lind=1)
plots.update([epoch, avg_loss_g], rind=1, cind=2, lind=1)
plots.save()
def validate(epoch):
epoch_loss_g = 0
total_psnr = 0
for i, batch in enumerate(validation_data_loader, 1):
if i > max_val_iter:
break
target = Variable(batch[0])
input = Variable(batch[1])
if cuda:
target = target.cuda()
input = input.cuda()
# Train Generator
reconstructed = gennet(input)
content_loss = content_criterion(reconstructed, target)
loss = content_loss
epoch_loss_g += loss.data[0]
total_psnr += psnr(loss.data[0])
print("Epoch[{}]({}/{}): Loss(G): {:.4f}".format(
epoch, i, max_val_iter, loss.data[0]))
avg_loss_g = epoch_loss_g / max_val_iter
avg_psnr = total_psnr / max_val_iter
print("Epoch {} : Avg Loss(G): {:.4f}".format(epoch, avg_loss_g))
plots.update([epoch, avg_psnr], rind=1, cind=1, lind=2)
plots.update([epoch, avg_loss_g], rind=1, cind=2, lind=2)
plots.save()
def validate(epoch):
epoch_loss_g = 0
epoch_loss_d = 0
total_psnr = 0
for i, batch in enumerate(validation_data_loader, 1):
real_batch_size = batch[0].size(0)
label.data.resize_(real_batch_size)
target = Variable(batch[0])
input = Variable(batch[1])
if cuda:
target = target.cuda()
input = input.cuda()
# Forward pass for Generator
reconstructed = gennet(input)
feature_r = vgg(reconstructed)
feature_t = vgg(input)
content_loss = content_criterion(feature_r, feature_t.detach())
label.data.fill_(real_label)
adversarial_loss = adversarial_criterion(disnet(reconstructed),
label)
loss = content_loss + 1e-3 * adversarial_loss
# Forward pass for Discriminator
reconstructed = gennet(input)
label.data.fill_(fake_label)
fake_loss = adversarial_criterion(disnet(reconstructed.detach()),
label)
label.data.fill_(real_label)
real_loss = adversarial_criterion(disnet(target), label)
epoch_loss_g += loss.data[0]
epoch_loss_d += fake_loss.data[0] + real_loss.data[0]
total_psnr += psnr(
content_criterion(reconstructed, target).data[0])
print("Epoch[{}]({}/{}): Loss(G): {:.4f} Loss(D): {:.4f}".format(
epoch, i, len(training_data_loader), loss.data[0],
fake_loss.data[0] + real_loss.data[0]))
avg_loss_g = epoch_loss_g / len(training_data_loader)
avg_loss_d = epoch_loss_d / len(training_data_loader)
avg_psnr = total_psnr / len(training_data_loader)
print("Epoch {} : Avg Loss(G): {:.4f} Avg Loss(D): {:.4f}".format(
epoch, avg_loss_g, avg_loss_d))
plots.update([epoch, avg_psnr], rind=1, cind=1, lind=2)
plots.update([epoch, avg_loss_g], rind=1, cind=2, lind=3)
plots.update([epoch, avg_loss_d], rind=1, cind=2, lind=4)
plots.save()
def validate(epoch):
total_loss = 0
total_psnr = 0
for batch in validation_data_loader:
target = Variable(batch[0])
input = Variable(batch[1])
if cuda:
target = target.cuda()
input = input.cuda()
mse = criterion(net(input), target)
total_loss += mse.data[0]
total_psnr += psnr(mse.data[0])
avg_loss = total_loss / len(validation_data_loader)
avg_psnr = total_psnr / len(validation_data_loader)
print("===> Avg. PSNR: {:.4f} dB".format(avg_psnr))
plots.update([epoch, avg_psnr], rind=1, cind=1, lind=2)
plots.update([epoch, avg_loss], rind=1, cind=2, lind=2)
plots.save()
def run():
image_seq = []
mat_data = sio.loadmat('./data/toy31_cassi.mat')
im_orig = mat_data['orig'] / 255.
im_orig = torch.from_numpy(im_orig).type(torch.float32).to(device)
image_m, image_n, image_c = im_orig.shape
mask = torch.from_numpy(mat_data['mask'].astype(np.float32)).to(device)
shape = im_orig.shape
y = mat_data['meas'] / 255.
y = torch.from_numpy(y).type(torch.float32).to(device)
# data missing and noise
# y = y + level * torch.randn_like(y)
# index_rand = np.random.rand(*list(y.shape))
# index_y = np.argwhere(index_rand < 0.05)
# y[index_y[:,0], index_y[:,1]] = 0
x = y.unsqueeze(2).expand_as(mask) * mask
mask_sum = torch.sum(mask**2, dim=2)
mask_sum[mask_sum == 0] = 1
flag = True
y1 = torch.zeros_like(y, dtype=torch.float32, device=device)
sigma = 5. / 255
for i in tqdm(range(100)):
if i == 20: flag = False
if i < 20: sigma = 50. / 255
elif i < 30: sigma = 25. / 255
elif i < 40: sigma = 12. / 255
else: sigma = 6. / 255
yb = torch.sum(mask * x, dim=2)
# no Acceleration
# temp = (y - yb) / mask_sum
# x = x + 1 * (temp.unsqueeze(2).expand_as(mask) * mask)
y1 = y1 + (y - yb)
temp = (y1 - yb) / mask_sum
x = x + 1 * (temp.unsqueeze(2).expand_as(mask) * mask)
x = ffdnet_denosing(x, sigma, flag).clamp(0, 1)
image_seq.append(x[:, :, 0].clamp(0., 1.).cpu().numpy())
# save_ani(image_res, filename='ffd_HSI.mp4', fps=10)
x.clamp_(0., 1.)
psnr_ = [psnr(x[..., kv], im_orig[..., kv]) for kv in range(image_c)]
ssim_ = [ssim(x[..., kv], im_orig[..., kv]) for kv in range(image_c)]
return np.mean(psnr_), np.mean(ssim_)
def ApertureWisePSNR(Groundtruth, Reconstruction):
"""
Calculate the PSNR value for each sub-aperture image of the
input reconstructed light field.
:param Groundtruth: input groundtruth light field
:param Reconstruction: input reconstruced light field
:return: aperture-wise PSNR values
"""
h, w, s, t = Groundtruth.shape[:4]
PSNRs = np.zeros([s, t])
for i in range(s):
for j in range(t):
gtimg = Groundtruth[:, :, i, j, ...]
gtimg = np.squeeze(gtimg)
recons = Reconstruction[:, :, i, j, ...]
recons = np.squeeze(recons)
PSNRs[i, j] = psnr(gtimg, recons)
return PSNRs
def train_epoch(epoch):
epoch_loss = 0
total_psnr = 0
for i, batch in enumerate(training_data_loader, 1):
target = Variable(batch[0])
input = Variable(batch[1])
if cuda:
target = target.cuda()
input = input.cuda()
optimizer.zero_grad()
loss = criterion(net(input), target)
epoch_loss += loss.data[0]
total_psnr += psnr(loss.data[0])
loss.backward()
optimizer.step()
print("Epoch[{}]({}/{}): Loss: {:.4f}".format(
epoch, i, len(training_data_loader), loss.data[0]))
avg_loss = epoch_loss / len(training_data_loader)
avg_psnr = total_psnr / len(training_data_loader)
print("Epoch {} : Avg Loss: {:.4f}".format(epoch, avg_loss))
plots.update([epoch, avg_psnr], rind=1, cind=1, lind=1)
plots.update([epoch, avg_loss], rind=1, cind=2, lind=1)
plots.save()
def pnp_fbs_csmri(denoise_func, im_orig, mask, noises, **opts):
alpha = opts.get('alpha', 0.4)
maxitr = opts.get('maxitr', 100)
verbose = opts.get('verbose', 1)
sigma = opts.get('sigma', 5)
""" Initialization. """
m, n = im_orig.shape
index = np.nonzero(mask)
y = np.fft.fft2(im_orig) * mask + noises # observed value
x_init = np.fft.ifft2(y) # zero fill
print(psnr(x_init, im_orig))
x = np.copy(x_init)
""" Main loop. """
for i in range(maxitr):
xold = np.copy(x)
""" Update variables. """
res = np.fft.fft2(x) * mask
index = np.nonzero(mask)
res[index] = res[index] - y[index]
x = x - alpha * np.fft.ifft2(res)
# x = np.real( x )
x = np.absolute(x)
""" Monitoring. """
# psnr
if verbose:
print("i: {}, \t psnr: {}"\
.format(i+1, psnr(x,im_orig)))
xout = np.copy(x)
""" Denoising step. """
xtilde = np.copy(x)
mintmp = np.min(xtilde)
maxtmp = np.max(xtilde)
xtilde = (xtilde - mintmp) / (maxtmp - mintmp)
# the reason for the following scaling:
# our denoisers are trained with "normalized images + noise"
# so the scale should be 1 + O(sigma)
scale_range = 1.0 + sigma / 255.0 / 2.0
scale_shift = (1 - scale_range) / 2.0
xtilde = xtilde * scale_range + scale_shift
# pytorch denoising model
xtilde_torch = np.reshape(xtilde, (1, 1, m, n))
xtilde_torch = torch.from_numpy(xtilde_torch).type(
torch.FloatTensor).cuda()
r = denoise_func(xtilde_torch).cpu().numpy()
r = np.reshape(r, (m, n))
x = xtilde - r
# scale and shift the denoised image back
x = (x - scale_shift) / scale_range
x = x * (maxtmp - mintmp) + mintmp
return xout
def train_epoch(epoch):
epoch_loss_g = 0
epoch_loss_d = 0
total_psnr = 0
for i, batch in enumerate(training_data_loader, 1):
real_batch_size = batch[0].size(0)
label.data.resize_(real_batch_size)
target = Variable(batch[0])
input = Variable(batch[1])
if cuda:
target = target.cuda()
input = input.cuda()
# Train Generator
optimizerG.zero_grad()
reconstructed = gennet(input)
feature_r = vgg(reconstructed)
feature_t = vgg(target)
content_loss = content_criterion(feature_r, feature_t.detach())
# use real label for log(G(D)) instead of 1-log(G(D))
label.data.fill_(real_label)
pred = disnet(reconstructed)
adversarial_loss = adversarial_criterion(pred, label)
loss = content_loss + 1e-3 * adversarial_loss
loss.backward()
optimizerG.step()
# Train Discriminator
optimizerD.zero_grad()
reconstructed = gennet(input)
label.data.fill_(fake_label)
fake_loss = adversarial_criterion(disnet(reconstructed.detach()),
label)
fake_loss.backward()
label.data.fill_(real_label)
real_loss = adversarial_criterion(disnet(target), label)
real_loss.backward()
optimizerD.step()
epoch_loss_g += loss.data[0]
epoch_loss_d += fake_loss.data[0] + real_loss.data[0]
total_psnr += psnr(
content_criterion(reconstructed, target).data[0])
print("Epoch[{}]({}/{}): Loss(G): {:.4f} Loss(D): {:.4f}".format(
epoch, i, len(training_data_loader), loss.data[0],
fake_loss.data[0] + real_loss.data[0]))
avg_loss_g = epoch_loss_g / len(training_data_loader)
avg_loss_d = epoch_loss_d / len(training_data_loader)
avg_psnr = total_psnr / len(training_data_loader)
print("Epoch {} : Avg Loss(G): {:.4f} Avg Loss(D): {:.4f}".format(
epoch, avg_loss_g, avg_loss_d))
plots.update([epoch, avg_psnr], rind=1, cind=1, lind=1)
plots.update([epoch, avg_loss_g], rind=1, cind=2, lind=1)
plots.update([epoch, avg_loss_d], rind=1, cind=2, lind=2)
plots.save()
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
with torch.no_grad():
preds = model(inputs).clamp(0.0, 1.0)
'''
Adding image grids slows dows the execution. Uncomment this block if you wish to log output images in Tensorboard
'''
# grid_inputs = torchvision.utils.make_grid(inputs)
# writer.add_image('Input LR',grid_inputs)
# grid_outputs = torchvision.utils.make_grid(preds)
# writer.add_image('Output HR',grid_outputs)
# grid_gt = torchvision.utils.make_grid(labels)
# writer.add_image('Ground Truth HR',grid_gt)
psnr1 = psnr(inputs.squeeze(0), preds.squeeze(0))
val_psnr_list += psnr1
writer.add_scalar("val_psnr", psnr1)
print('validation psnr: {} for epoch: {}'.format(
val_psnr_list / len(val_dataset), epoch))
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'val_psnr': psnr1,
}, os.path.join('saved_weights/weighted',
'epoch_{}.pth'.format(epoch)))
outputs = model(inputs)
loss = criterion(outputs, labels).to(device)
loss.backward()
optimizer.step
writer.add_scalar('loss', loss.item())
running_loss += loss.item() * inputs.size(0)
print('Epoch {}, Iteration: {}, Loss: {}'.format(
epoch, iteration, loss.item()))
'''
calculating psnr,msim for every iteration is expensive. Can calculate for every 500 iterations
'''
if (iteration % 100 == 0):
psnr_score = psnr(inputs, outputs)
writer.add_scalar('psnr_train', psnr_score)
ssim_score = utils.pytorch_ssim.ssim(
inputs, outputs) #calculated in batches of batch_size
writer.add_scalar('ssim_train', ssim_score.item())
if (phase == 'validation'):
if (iteration % 50 == 0):
psnr_score = psnr(inputs, outputs)
writer.add_scalar('psnr_validation', psnr_score)
ssim_score = utils.pytorch_ssim.ssim(
inputs, outputs) #calculated in batches of batch_size
writer.add_scalar('ssim_validation', ssim_score.item())
if (iteration % 200 == 0):
grid_inputs = torchvision.utils.make_grid(inputs)
writer.add_image('Input LR', grid_inputs)
def train(patch_width=PATCH_WIDTH,
patch_height=PATCH_HEIGHT,
epoch_num=EPOCH_NUM,
batch_size=BATCH_SIZE,
learning_rate=LEARNING_RATE,
learning_rate_decay_steps=LEARNING_RATE_DECAY_STEPS,
learning_rate_decay_rate=LEARNING_RATE_DECAY_RATE,
dsp_itv=DSP_ITV,
ckpt_path=ckpt_path,
save_best=save_best,
early_stop=early_stop):
'''
Difine the tensorflow graph, execute training and validation.
'''
#------------------------------define the graph---------------------------
Y = tf.placeholder(tf.float32, [None, None, None, 1], name='Y')
S0 = tf.placeholder(tf.float32, [None, None, None, 1], name='S0')
DoLP = tf.placeholder(tf.float32, [None, None, None, 1], name='DoLP')
AoP = tf.placeholder(tf.float32, [None, None, None, 1], name='AoP')
# Para = tf.placeholder(tf.float32, [None, None, None, 3])#define tensors of input and label
# DoLP_hat= ForkNet(Y)
S0_hat, DoLP_hat, AoP_hat = ForkNet(Y, padding='SAME')
loss = LOSS(S0_hat, S0, DoLP_hat, DoLP, AoP_hat, AoP)
tf.add_to_collection('S0_hat', S0_hat)
tf.add_to_collection('DoLP_hat', DoLP_hat)
tf.add_to_collection('AoP_hat', AoP_hat)
global_step = tf.Variable(0, name='global_step', trainable=False)
decayed_learning_rate = tf.train.exponential_decay(
learning_rate,
global_step,
learning_rate_decay_steps,
learning_rate_decay_rate,
staircase=True) # lerning rate decay
train_step = tf.train.AdamOptimizer(decayed_learning_rate).minimize(
loss, global_step=global_step)
init = tf.global_variables_initializer()
saver = tf.train.Saver() #instantiate the saveer
train_steps, train_Y, train_label, val_steps, val_Y, val_label = load_data(
)
# val_s0 = val_para[:, :, :, :1]
# val_dolp = val_para[:, :, :, 1:2]
# val_aop = val_para[:, :, :, 2:]
# print(np.max(val_Y))
with tf.Session() as sess:
print('\nStart Training!')
sess.run(init)
min_loss = np.inf
wait = 0
psnr_record = []
total_S0_PSNR_BIC = 0
total_DoLP_PSNR_BIC = 0
total_AoP_PSNR_BIC = 0
for epoch in range(epoch_num):
#training set batch generator
train_generator = patch_batch_generator(train_Y,
train_label,
batch_size,
patch_width,
patch_height,
random_shuffle=True)
#test set batch generator
val_generator = patch_batch_generator(val_Y,
val_label,
batch_size,
patch_width,
patch_height,
random_shuffle=False,
augment=False)
print(
'=======================================Epoch:{}/{}======================================='
.format(epoch, epoch_num))
# training
total_train_loss = 0
for step in range(train_steps):
(Input_batch_train, Para_batch_train) = next(train_generator)
Y_batch_train = Input_batch_train[:, :, :, :1]
S0_batch_train = Para_batch_train[:, :, :, :1]
DoLP_batch_train = Para_batch_train[:, :, :, 1:2]
AoP_batch_train = Para_batch_train[:, :, :, 2:]
sess.run(train_step,
feed_dict={
Y: Y_batch_train,
S0: S0_batch_train,
DoLP: DoLP_batch_train,
AoP: AoP_batch_train
})
train_loss = sess.run(loss,
feed_dict={
Y: Y_batch_train,
S0: S0_batch_train,
DoLP: DoLP_batch_train,
AoP: AoP_batch_train
})
total_train_loss += train_loss
if step % dsp_itv == 0:
rate = float(step + 1) / float(train_steps)
rate_num = int(rate * 100)
arrow = 0 if step + 1 == train_steps else 1
r = '\rStep:%d/%d [%s%s%s]%d%% --- Training loss:%f' % \
(step + 1, train_steps, '■' * rate_num, '▶' * arrow, '-' * (100 - rate_num - arrow), rate * 100, train_loss)
print(r)
# validation
total_val_loss = 0
total_S0_PSNR = 0
total_DoLP_PSNR = 0
total_AoP_PSNR = 0
for step in range(val_steps):
(Input_batch_val, Para_batch_val) = next(val_generator)
Y_batch_val = Input_batch_val[:, :, :, :1]
BIC_batch_val = Input_batch_val[:, :, :, 1:]
S0_batch_val = Para_batch_val[:, :, :, :1]
DoLP_batch_val = Para_batch_val[:, :, :, 1:2]
AoP_batch_val = Para_batch_val[:, :, :, 2:]
total_val_loss += sess.run(loss,
feed_dict={
Y: Y_batch_val,
S0: S0_batch_val,
DoLP: DoLP_batch_val,
AoP: AoP_batch_val
})
S0_hat_val, DoLP_hat_val, AoP_hat_val = sess.run(
[S0_hat, DoLP_hat, AoP_hat], feed_dict={Y: Y_batch_val})
#limit the value
S0_hat_val = np.clip(S0_hat_val, 0, 2)
DoLP_hat_val = np.clip(DoLP_hat_val, 0, 1)
# AoP_hat_val = np.clip(AoP_hat_val, 0, math.pi)
# DoLP_hat_val = Normalize(DoLP_hat_val, 0, 1)
total_S0_PSNR += psnr(S0_batch_val[:, :, :, 0],
S0_hat_val[:, :, :, 0], 2)
total_DoLP_PSNR += psnr(DoLP_batch_val[:, :, :, 0],
DoLP_hat_val[:, :, :, 0], 1)
total_AoP_PSNR += psnr(AoP_batch_val[:, :, :, 0],
AoP_hat_val[:, :, :, 0], math.pi / 2.)
# for b in range(AoP_batch_val.shape[0]):
# total_AoP_PSNR += compare_ssim(np.float32(AoP_batch_val[b, :, :, 0]), np.float32(AoP_hat_val[b, :, :, 0]), data_range=math.pi / 2.)
if epoch == 0:
# print('max:', max(val_bic[0,6:-6,6:-6,0]))
S0_BIC = (BIC_batch_val[:, :, :, 0] +
BIC_batch_val[:, :, :, 1] +
BIC_batch_val[:, :, :, 2] +
BIC_batch_val[:, :, :, 3]) / 2.
DoLP_BIC = dolp(BIC_batch_val[:, :, :, 0],
BIC_batch_val[:, :, :, 1],
BIC_batch_val[:, :, :,
2], BIC_batch_val[:, :, :,
3])
AoP_BIC = aop(
BIC_batch_val[:, :, :, 0], BIC_batch_val[:, :, :, 1],
BIC_batch_val[:, :, :, 2], BIC_batch_val[:, :, :, 3]
) + math.pi / 4. #avoid the minus number
total_S0_PSNR_BIC += psnr(S0_batch_val[:, :, :, 0], S0_BIC,
2)
total_DoLP_PSNR_BIC += psnr(DoLP_batch_val[:, :, :, 0],
DoLP_BIC, 1)
total_AoP_PSNR_BIC += psnr(AoP_batch_val[:, :, :, 0],
AoP_BIC, math.pi / 2.)
# for b in range(AoP_batch_val.shape[0]):
# total_AoP_PSNR_BIC += compare_ssim(np.float32(AoP_batch_val[b, :, :, 0]), np.float32(AoP_BIC[b]), data_range=math.pi / 2.)
print(
'========================================Validation======================================='
+ '\nTraining loss: %.5f' % (total_train_loss / train_steps) +
'\nValidation loss: %.5f' % (total_val_loss / val_steps) +
'\n ————————————————————————————————————————————————————————————————————————————————'
+
# '\n| PSNR of I_0 using PDCNN: %.5f | PSNR of I_0 using BICUBIC: %.5f |' % (total_X_0_PSNR/val_steps, 32.872) +
# '\n| PSNR of I_45 using PDCNN: %.5f | PSNR of I_45 using BICUBIC: %.5f |' % (total_X_45_PSNR/val_steps, 32.973) +
# '\n| PSNR of I_90 using PDCNN: %.5f | PSNR of I_90 using BICUBIC: %.5f |' % (total_X_90_PSNR/val_steps, 33.008) +
# '\n| PSNR of I_135 using PDCNN: %.5f | PSNR of I_135 using BICUBIC: %.5f |' % (total_X_135_PSNR/val_steps, 32.923) +
'\n| PSNR of S_0 using SRCNN: %.5f | PSNR of S_0 using BICUBIC: %.5f |'
% (total_S0_PSNR / val_steps, total_S0_PSNR_BIC / val_steps) +
'\n| PSNR of DoLP using SRCNN: %.5f | PSNR of DoLP using BICUBIC: %.5f |'
% (total_DoLP_PSNR / val_steps,
total_DoLP_PSNR_BIC / val_steps) +
'\n| PSNR of AoP using SRCNN: %.5f | PSNR of AoP using BICUBIC: %.5f |'
%
(total_AoP_PSNR / val_steps, total_AoP_PSNR_BIC / val_steps) +
'\n ————————————————————————————————————————————————————————————————————————————————'
)
psnr_record.append([
total_S0_PSNR / val_steps, total_DoLP_PSNR / val_steps,
total_AoP_PSNR / val_steps
])
if save_best or early_stop:
if metrics == 'validation loss':
current_loss = total_val_loss / val_steps
elif metrics == 'training loss':
current_loss = total_train_loss / train_steps
if current_loss < min_loss:
print('Validation loss decreased from %.5f to %.5f' %
(min_loss, current_loss))
min_loss = current_loss
if save_best:
saver.save(sess, ckpt_path)
print("Model saved in file: %s" % ckpt_path)
if early_stop:
wait = 0
else:
print('Validation loss did not decreased.')
if early_stop:
wait += 1
if wait > patient:
print('Early stop!')
break
if not save_best:
saver.save(sess, ckpt_path)
print("Model saved in file: %s" % ckpt_path)
with open(csv_path, 'w') as csv_file:
writer = csv.writer(csv_file)
writer.writerows(psnr_record)
def pnp_admm_photon_imaging(b, denoiser, im_true, **opts):
"""
Parameters:
:b - the observation in single photon imaging.
:denoiser - the Gaussian denoiser used in Plug-and-Play ADMM.
:im_true - the clean image used to monitor PSNR.
:opts - the kwargs for hyperparameters in Plug-and-Play ADMM.
:K - the parameter in single photo imaging.
:lam - the value of 1/alpha.
:rho - TODO
:maxitr - the max number of iterations.
:verbose - a flag that enables/disables info printing.
- NOTE: if `peak` and `M` options exist in `opts`, then the
`clean` image is the scaled version of the original image.
:beta - the prior weight parameter.
:step - TODO
"""
""" Process parameters. """
K = opts.get('K', 8)
lam = opts.get('lam', 15.0)
rho = opts.get('rho', 100.0) # rho = 1.0 / alpha
maxitr = opts.get('maxitr', 50)
data_range = opts.get('data_range', 1.0)
verbose = opts.get('verbose', 1)
step = opts.get('step', 1.0)
beta = opts.get('beta', 1.0)
""" Initialization. """
K1 = blockfunc(b, (K, K), np.sum)
x = K1 / K**2
u = np.zeros_like(x, dtype=np.float64)
v = np.copy(x)
m, n = x.shape
sigma = np.sqrt(lam / rho)
""" Main loop. """
for i in range(maxitr):
x_old = x
u_old = u
v_old = v
""" Inverse step. """
x = inverse_step(u, v, K1, K, rho)
""" Denoising step. """
vtilde = x + u
# scale vtilde to be in range of [0,1]
mintmp = 0.0
maxtmp = np.max(vtilde)
vtilde = (vtilde - mintmp) / (maxtmp - mintmp)
trans_sigma = sigma / (maxtmp - mintmp)
# then set data range to [0.15, 0.85] to avoid clipping of extreme values
scale_range = 0.4
scale_shift = (1 - scale_range) / 2.0
vtilde = vtilde * scale_range + scale_shift
trans_sigma = trans_sigma * scale_range
# pytorch denoising model
vtilde_torch = np.reshape(vtilde, (1, 1, m, n))
vtilde_torch = torch.from_numpy(vtilde_torch).type(
torch.FloatTensor).cuda()
r = denoiser(vtilde_torch).cpu().numpy()
r = np.reshape(r, (m, n))
v = vtilde - r
# scale and shift the denoised image back
v = (v - scale_shift) / scale_range
v = v * (maxtmp - mintmp) + mintmp
""" Update variables. """
u = u + x - v
""" Monitoring. """
# successive difference
dif = np.sqrt(np.sum(np.square(x - x_old)))
dif_denom = np.sqrt(np.sum(np.square(x_old)))
# psnr
if verbose:
print("i: {}, \t successive difference: {}, \t psnr: {} \t"\
.format(i+1, dif/dif_denom, psnr(im_true, x)))
return x
15, 2.0, 100
) # the arguments are default sigma, default alpha and default max iteration.
# ---- load the model ----
model = load_model(opt.model_type, opt.sigma)
with torch.no_grad():
# ---- load the ground truth ----
im_orig = cv2.imread('Demo_mat/CS_MRI/Brain.jpg', 0) / 255.0
# ---- load mask matrix ----
mat = sio.loadmat('Demo_mat/CS_MRI/Q_Random30.mat')
mask = mat.get('Q1').astype(np.float64)
# ---- load noises -----
noises = sio.loadmat('Demo_mat/CS_MRI/noises.mat')
noises = noises.get('noises').astype(np.complex128) * 3.0
# ---- set options -----
opts = dict(sigma=opt.sigma,
alpha=opt.alpha,
maxitr=opt.maxitr,
verbose=opt.verbose)
# ---- plug and play !!! -----
out = pnp_admm_csmri(model, im_orig, mask, noises, **opts)
# ---- print result -----
print('Plug-and-Play PNSR: ', psnr(out, im_orig))
DoLP_true = dolp(origin_img[i, :, :, 0], origin_img[i, :, :, 1], origin_img[i, :, :, 2], origin_img[i, :, :, 3])
AoP_true = aop(origin_img[i, :, :, 0], origin_img[i, :, :, 1], origin_img[i, :, :, 2], origin_img[i, :, :, 3]) + math.pi/4.
origin_s0[i] = S0_true
origin_dolp[i] = DoLP_true
origin_aop[i] = AoP_true
#the dolp of bic images
S0_BIC = 1 / 2 * (bic_img[i, :, :, 0] + bic_img[i, :, :, 1] + bic_img[i, :, :, 2] + bic_img[i,:, :, 3])
DoLP_BIC = dolp(bic_img[i, :, :, 0], bic_img[i, :, :, 1], bic_img[i, :, :, 2], bic_img[i, :, :, 3])
AoP_BIC = aop(bic_img[i, :, :, 0], bic_img[i, :, :, 1], bic_img[i, :, :, 2], bic_img[i, :, :, 3]) + math.pi / 4.
bic_s0[i] = S0_BIC
bic_dolp[i] = DoLP_BIC
bic_aop[i] = AoP_BIC
# Calculate the PSNR of S0, DoLP and AoP obtained through PDCNN method
total_S0_PSNR[i]= psnr(S0_true, S0_hat_test, 2)
total_DoLP_PSNR[i] = psnr(DoLP_true, DoLP_hat_test, 1)
total_AoP_PSNR[i] = psnr(AoP_true, AoP_hat_test, math.pi/2.)
# Calculate the PSNR of S0, DoLP and AoP obtained through BICUBIC method
total_S0_PSNR_BIC[i] = psnr(S0_true, S0_BIC, 2)
total_DoLP_PSNR_BIC[i] = psnr(DoLP_true, DoLP_BIC, 1)
total_AoP_PSNR_BIC[i] = psnr(AoP_true, AoP_BIC, math.pi/2.)
# show the progress bar
view_bar(i, IMG_NUM)
print('\n========================================Testing=======================================' +
'\n ————————————————————————————————————————————————————————————————————————————————' +
'\n| PSNR of S_0 using SRCNN: %.5f | PSNR of S_0 using BICUBIC: %.5f |' % (np.mean(total_S0_PSNR), np.mean(total_S0_PSNR_BIC)) +
'\n| PSNR of DoLP using PDCNN: %.5f | PSNR of DoLP using BICUBIC: %.5f |' % (np.mean(total_DoLP_PSNR), np.mean(total_DoLP_PSNR_BIC)) +
def pnp_admm_csmri(model, im_orig, mask, noises, **opts):
alpha = opts.get('alpha', 2.0)
maxitr = opts.get('maxitr', 100)
verbose = opts.get('verbose', 1)
sigma = opts.get('sigma', 5)
""" Initialization. """
m, n = im_orig.shape
index = np.nonzero(mask)
y = np.fft.fft2(im_orig) * mask + noises # observed value
x_init = np.fft.ifft2(y) # zero fill
print('zero-fill PSNR:', psnr(x_init, im_orig))
x = np.absolute(np.copy(x_init))
v = np.copy(x)
u = np.zeros((m, n), dtype=np.float64)
""" Main loop. """
for i in range(maxitr):
xold = np.copy(x)
vold = np.copy(v)
uold = np.copy(u)
""" Update variables. """
vtilde = np.copy(x + u)
vf = np.fft.fft2(vtilde)
La2 = 1.0 / 2.0 / alpha
vf[index] = (La2 * vf[index] + y[index]) / (1.0 + La2)
v = np.real(np.fft.ifft2(vf))
""" Denoising step. """
xtilde = np.copy(2 * v - xold - uold)
mintmp = np.min(xtilde)
maxtmp = np.max(xtilde)
xtilde = (xtilde - mintmp) / (maxtmp - mintmp)
# the reason for the following scaling:
# our denoisers are trained with "normalized images + noise"
# so the scale should be 1 + O(sigma)
scale_range = 1.0 + sigma / 255.0 / 2.0
scale_shift = (1 - scale_range) / 2.0
xtilde = xtilde * scale_range + scale_shift
# pytorch denoising model
xtilde_torch = np.reshape(xtilde, (1, 1, m, n))
xtilde_torch = torch.from_numpy(xtilde_torch).type(
torch.FloatTensor).cuda()
r = model(xtilde_torch).cpu().numpy()
r = np.reshape(r, (m, n))
x = xtilde - r
# scale and shift the denoised image back
x = (x - scale_shift) / scale_range
x = x * (maxtmp - mintmp) + mintmp
""" Update variables. """
u = uold + xold - v
""" Monitoring. """
if verbose:
print("i: {}, \t psnr: {}"\
.format(i+1, psnr(x,im_orig)))
return x
def run():
mat_data = sio.loadmat('./data/toy31_cassi.mat')
im_orig = mat_data['orig'] / 255
im_orig = torch.from_numpy(im_orig).type(torch.float32).to(device)
image_m, image_n, image_c = im_orig.shape
# image_seq = []
# ---- load mask matrix ----
mask = torch.from_numpy(mat_data['mask'].astype(np.float32)).to(device)
y = mat_data['meas'] / 255
y = torch.from_numpy(y).type(torch.float32).to(device)
# data missing and noise
# y = y + level * torch.randn_like(y)
# index_rand = np.random.rand(*list(y.shape))
# index_y = np.argwhere(index_rand < 0.05)
# y[index_y[:,0], index_y[:,1]] = 0
x = y.unsqueeze(2).expand_as(mask) * mask
mask_sum = torch.sum(mask**2, dim=2)
mask_sum[mask_sum == 0] = 1
flag = True
y1 = torch.zeros_like(y, dtype=torch.float32, device=device)
sigma_ = 50 / 255
for i in tqdm(range(100)):
if i == 20: flag = False
yb = torch.sum(mask * x, dim=2)
# no Acceleration
# temp = (y - yb) / (mask_sum)
# x = x + 1 * (temp.unsqueeze(2).expand_as(mask) * mask)
y1 = y1 + (y - yb)
temp = (y1 - yb) / mask_sum
x = x + 1 * (temp.unsqueeze(2).expand_as(mask) * mask)
if i < 20:
x = ffdnet_denosing(x, 50. / 255, flag)
else:
ffdnet_hypara_list = [100., 80., 60., 40., 20., 10., 5.]
ffdnet_num = len(ffdnet_hypara_list)
tv_hypara_list = [10, 0.01]
tv_num = len(tv_hypara_list)
ffdnet_list = [
ffdnet_denosing(x, level / 255., flag).clamp(0, 1)
for level in ffdnet_hypara_list
]
tv_list = [
TV_denoising(x, level, 5).clamp(0, 1)
for level in tv_hypara_list
]
ffdnet_mat = np.stack([
x_ele[:, :, :].cpu().numpy().reshape(-1).astype(np.float64)
for x_ele in ffdnet_list
],
axis=0)
tv_mat = np.stack([
x_ele[:, :, :].cpu().numpy().reshape(-1).astype(np.float64)
for x_ele in tv_list
],
axis=0)
w = cp.Variable(ffdnet_num + tv_num)
P = np.zeros((ffdnet_num + tv_num, ffdnet_num + tv_num))
P[:ffdnet_num, :ffdnet_num] = ffdnet_mat @ ffdnet_mat.T
P[:ffdnet_num, ffdnet_num:] = -ffdnet_mat @ tv_mat.T
P[ffdnet_num:, :ffdnet_num] = -tv_mat @ ffdnet_mat.T
P[ffdnet_num:, ffdnet_num:] = tv_mat @ tv_mat.T
one_vector_ffdnet = np.ones((1, ffdnet_num))
one_vector_tv = np.ones((1, tv_num))
objective = cp.quad_form(w, P)
problem = cp.Problem(cp.Minimize(objective), [
one_vector_ffdnet @ w[:ffdnet_num] == 1,
one_vector_tv @ w[ffdnet_num:] == 1, w >= 0
])
problem.solve()
w_value = w.value
x_ffdnet, x_tv = 0, 0
for idx in range(ffdnet_num):
x_ffdnet += w_value[idx] * ffdnet_list[idx]
for idx in range(tv_num):
x_tv += w_value[idx + ffdnet_num] * tv_list[idx]
x = 0.5 * (x_ffdnet + x_tv)
# image_seq.append(x[...,0])
x.clamp_(0, 1)
# fps = 10
# save_ani(image_seq, filename='HSI.mp4', fps=fps)
psnr_ = [psnr(x[..., kv], im_orig[..., kv]) for kv in range(image_c)]
ssim_ = [ssim(x[..., kv], im_orig[..., kv]) for kv in range(image_c)]
return np.mean(psnr_), np.mean(ssim_)