def validation(model_path, file_data):
if torch.cuda.is_available():
model = NET.VDN_NET(in_channels=3, depth_snet=5).cuda()
model.load_state_dict(torch.load(model_path))
else:
model = NET.VDN_NET(in_channels=3, depth_snet=5)
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()
avg_psnr_validation = 0.0
avg_ssim_validation = 0.0
obj_data = gd.ValidationBenchmark(h5_file_=file_data)
if torch.cuda.is_available():
model.cuda()
torch.backends.cudnn.benchmark = True
for idx in range(obj_data.__len__()):
noisy, image = obj_data.__getitem__(idx)
ch, ht, wt = noisy.shape
noisy = noisy.view(1, ch, ht, wt).cuda()
image = image.cuda()
model_out, _ = model(noisy)
noise = noisy - model_out[:, :ch, ].detach().data
clean_img_pred = noise.view(ch, ht, wt).permute(1, 2, 0).clamp(0, 1)
image = image.view(ch, ht, wt).permute(1, 2, 0)
avg_psnr_validation += psnr(image * 255, clean_img_pred * 255)
avg_ssim_validation += compare_ssim(img_as_ubyte(image.cpu().numpy()),
img_as_ubyte(clean_img_pred.cpu().numpy()),
win_size=11, data_range=255, multichannel=True, gaussian_weights=True)
print("average validation PSNR = ", avg_psnr_validation / obj_data.__len__())
print("average validation SSIM = ", avg_ssim_validation / obj_data.__len__())
def testCase1():
qp = 15
step = 16
im = plt.imread("E:/liumangxuxu/code/PyCodec/modules/lena2.tif").astype(
int)
print(im.shape)
residual_1D, mode_1D = predictImage(im, qp, step)
i_im = inversePredictImage(residual_1D, mode_1D, qp, im.shape[0],
im.shape[1], step)
plt.figure()
plt.imshow(i_im, cmap='gray')
plt.title("Inverse image")
logging.debug("Inverse Image PSNR:")
logging.debug(tools.psnr(im, i_im))
plt.show()
def processCheck():
# 4x4 block test data set
# N = 4
# S = np.array([[1, 2, 2, 0],
# [0, 1, 3, 1],
# [0, 1, 2, 1],
# [1, 2, 2, -1]])
# real 8x8 image test data
N = 8
S = np.array([[182, 196, 199, 201, 203, 201, 199, 173],
[175, 180, 176, 142, 148, 152, 148, 120],
[148, 118, 123, 115, 114, 107, 108, 107],
[115, 110, 110, 112, 105, 109, 101, 100],
[104, 106, 106, 102, 104, 95, 98, 105],
[99, 115, 131, 104, 118, 86, 87, 133],
[112, 154, 154, 107, 140, 97, 88, 151],
[145, 158, 178, 123, 132, 140, 138, 133]])
# random test data
# N = 8 # axis of 1D vector
# S = np.random.randint(0, 100, size=[N, N])
print("\nTest 2D %dx%d matrix:" % (N, N))
print(S)
# calculate DCT patterns by formula
temp, S_i = dct_detail(S, N)
diff = tools.psnr(S, np.round(S_i))
print("\nPSNR is: ", np.around(diff, 2))
# check the coefficients by Scipy
T_scipy = np.round(block2dct(S), 4)
print("\nthe expansion coefficients using SciPy:")
print(T_scipy)
#something wrong with below calculation
S_scipy = np.round(dct2block(T_scipy), 4)
print("\nInverse Result using SciPy:")
print(S_scipy)
def processWholeImage():
im = plt.imread("lena2.tif").astype(float)
print(im.shape)
BLOCK_WIDTH = 8
# using Scipy way
dct = Img2DctUsingScipy(im, BLOCK_WIDTH)
img_dct = Dct2ImgUsingScipy(dct, BLOCK_WIDTH)
# using my own calculation way
#dct, img_dct = ImgDctUsingDetail(im, BLOCK_WIDTH)
plt.figure()
plt.imshow(dct, cmap='gray', vmax=np.max(dct) * 0.01, vmin=0)
plt.title("8x8 DCTs of the image")
plt.figure()
plt.imshow(np.hstack((im, img_dct)), cmap='gray')
plt.title("Comparison between original and DCT compressed images")
diff = tools.psnr(im, img_dct)
print("\nPSNR is: ", np.around(diff, 2))
plt.show()
def simulated_noise_test(noise_type,
model_path,
plot_noise=False,
stats=False,
plot_images=False,
conf1=False):
test_sets = {"LIVE1": "*.bmp", "CBSD68": "*.png", "Set5": "*.bmp"}
test_paths = sorted(
gd.load_data(os.path.join("datasets", "test_data"), test_sets))
data_obj = gd.TestDataset(test_paths, noise_type=noise_type, iid=True)
if torch.cuda.is_available():
if conf1:
model = NET1.VDN_NET(in_channels=3, depth_snet=5).cuda()
else:
model = NET.VDN_NET(in_channels=3, depth_snet=5).cuda()
model.load_state_dict(torch.load(model_path))
else:
model = NET.VDN_NET(in_channels=3, depth_snet=5)
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()
PSNR = 0.0
SSIM = 0.0
for idx in range(data_obj.__len__()):
image, noisy = data_obj.__getitem__(idx)
if torch.cuda.is_available():
image, noisy = Variable(image.cuda()), Variable(noisy.cuda())
_, ch, ht, wt = noisy.shape
model_out, model_out2 = model(noisy)
alpha = model_out2[:, :ch, ]
beta = model_out2[:, ch:, ]
result = beta / (alpha + 1)
noise = noisy - model_out[:, :ch, ]
clean_img_pred = noise.view(ch, ht, wt).permute(1, 2, 0).clamp(0, 1)
image = image.view(ch, ht, wt).permute(1, 2, 0)
if stats:
PSNR += psnr(image * 255, clean_img_pred * 255)
SSIM += compare_ssim(img_as_ubyte(image.cpu().detach().numpy()),
img_as_ubyte(
clean_img_pred.cpu().detach().numpy()),
multichannel=True)
if plot_noise:
fig = plt.figure()
ax = plt.axes(projection='3d')
x = np.arange(0, ht, 1)
y = np.arange(0, wt, 1)
X, Y = np.meshgrid(x, y)
Z1 = np.exp((img_as_float(result[0, 1, ].detach().numpy())))
surf = ax.plot_surface(X[:100, :100],
Y[:100, :100],
Z1[:100, :100],
rstride=1,
cstride=1,
cmap="viridis",
antialiased=True,
edgecolor="none")
plt.savefig("simulation_results/approx_noise2_" + str(idx) +
".png")
plt.show()
if plot_images:
plt.subplots_adjust(wspace=0.2)
plt.subplot(131)
plt.imshow(image[100:250, 100:250, ])
plt.title('Groundtruth')
plt.subplot(132)
plt.imshow(noisy[0].permute(1, 2, 0)[100:250, 100:250, ])
plt.title('Noisy Image')
plt.subplot(133)
plt.imshow(
img_as_ubyte(clean_img_pred.detach().numpy())[100:250,
100:250, ])
plt.title('Denoised Image')
plt.savefig("simulation_results/denoised_simlation" + str(idx) +
".png")
plt.show()
print("average PSNR = ", PSNR / data_obj.__len__())
print("average SSIM = ", SSIM / data_obj.__len__())
def training_model_VDN(arg, milestones):
logging.basicConfig(filename=arg.log_path, level=logging.INFO) # log file
logging.info('Started')
if not os.path.exists(arg.model_path):
os.makedirs(arg.model_path)
# load training data
train_sets = {
"BSDS500": "*.jpg",
"Waterloo": "*.bmp",
"ImageNet": "*.JPEG"
}
if arg.confnet == 1:
model = VDN_NET_one(in_channels=arg.channels, depth_snet=arg.snet)
else:
model = VDN_NET(in_channels=arg.channels, depth_snet=arg.snet)
model = model.float()
train_paths = np.array(sorted(gd.load_data("datasets", train_sets)))
ran = bool(arguments.randomize)
clipping = bool(arguments.clipping)
data_obj = gd.TrainDataset(train_paths,
randomize=ran,
patch_size=arg.patch)
if torch.cuda.is_available():
model.cuda()
torch.backends.cudnn.benchmark = True
data = DataLoader(dataset=data_obj,
drop_last=True,
batch_size=arg.batch,
shuffle=True,
num_workers=arg.workers)
else:
data = DataLoader(dataset=data_obj,
drop_last=True,
batch_size=arg.batch,
shuffle=True)
# network parameters
epsilon = np.sqrt(5.0e-6)
p_window = 7
optimizer = optim.Adam(model.parameters(), lr=2e-4)
scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=arg.gamma)
if clipping:
gadient_clip_Dnet = 1000.0
gadient_clip_Snet = 50.0
Dnet_parameters = [
x for name, x in model.named_parameters()
if 'dnet' in name.lower()
]
Snet_parameters = [
x for name, x in model.named_parameters()
if 'snet' in name.lower()
]
print("Training model simulation now!")
for epoch in range(arg.epochs):
if clipping:
grad_D = 0.0
grad_S = 0.0
epoch_avg_loss = 0.0
mse_avg = 0.0
psnr_avg = 0.0
ssim_avg = 0.0
lr = optimizer.param_groups[0]['lr']
if lr < arg.learning:
print("reach min learning rate at epoch" + str(epoch))
for i, batch_data in enumerate(data):
if torch.cuda.is_available():
x_batch, y_batch, sigma_arr = Variable(
batch_data[0]).cuda(), Variable(
batch_data[1]).cuda(), Variable(batch_data[2]).cuda()
else:
x_batch, y_batch, sigma_arr = batch_data[0], batch_data[
1], batch_data[2]
optimizer.zero_grad()
out_D, out_s = model(y_batch)
loss, loglikelihood, kl_z, kl_sigma = loss_func.get_loss(
x_batch, y_batch, sigma_arr, p_window,
out_D[:, :arg.channels, :, :], out_D[:, arg.channels:, :, :],
out_s[:, :arg.channels, :, :], out_s[:, arg.channels:, :, :],
epsilon)
loss.backward()
if clipping:
full_grad_D = nn.utils.clip_grad_norm_(Dnet_parameters,
gadient_clip_Dnet)
full_grad_S = nn.utils.clip_grad_norm_(Snet_parameters,
gadient_clip_Snet)
grad_D = (grad_D * (i / (i + 1)) + full_grad_D / (i + 1))
grad_S = (grad_S * (i / (i + 1)) + full_grad_S / (i + 1))
optimizer.step()
epoch_avg_loss += loss.detach().item()
predicted_image = y_batch - out_D[:, :arg.channels, :, :].detach(
).data
predicted_image = predicted_image.clamp(0, 1)
mse = calc_MSE(predicted_image, x_batch)
mse_avg += mse
psnr_avg += psnr(predicted_image * 255, x_batch * 255)
ssim_avg += calculate_ssim(img_as_ubyte(
predicted_image.permute(2, 3, 1, 0).cpu().numpy()),
img_as_ubyte(
x_batch.permute(2, 3, 1,
0).cpu().numpy()),
multichannel=True)
if i == 0:
print("First ForwardPAss\n Loss: {}, MSE: {}".format(
loss.detach().item(), mse))
if (i + 1) % 100 == 0:
print("{} - Loss: {}, MSE:{}, epoch:{}".format(
i + 1, loss.item(), mse, epoch + 1))
if i >= 5000:
break
if clipping:
gadient_clip_Dnet = min(gadient_clip_Dnet, grad_D)
gadient_clip_Dnet = min(gadient_clip_Dnet, grad_S)
print("----------------------------------------------------------")
print(
"Epoch: {}, Avg MSE:{}, Avg Epoch Loss:{}, Avg PSNR:{}, Avg SSIM : {}, LR:{}"
.format(epoch + 1, mse_avg / (i + 1), epoch_avg_loss / (i + 1),
psnr_avg / (i + 1), ssim_avg / (i + 1), lr))
logging.info("av loss: {}, epoch: {}".format(epoch_avg_loss / (i + 1),
epoch + 1))
scheduler.step()
if epoch % arguments.epoch_save == 0:
torch.save(
model.state_dict(),
os.path.join(arg.model_path,
"model_" + str(epoch) + "_epochs.pth"))
print("saved model as" + arg.model_path)
print("Finished Training...\n Saving model now.....\n")
torch.save(model.state_dict(),
os.path.join(arg.model_path, "final_model.pth"))
print("saved model as" + os.path.join(arg.model_path, "final_model.pth"))
def training_benchmark(arg, milestones):
logging.basicConfig(filename=arg.log_path, level=logging.INFO) # log file
logging.info('Started')
if not os.path.exists(arg.model_path):
os.makedirs(arg.model_path)
model = VDN_NET(in_channels=arg.channels, depth_snet=arg.snet)
model = model.float()
clipping = bool(arguments.clipping)
# Load training data
obj_data = gd.TrainBenchmark(h5_file_=arg.train_data,
patch_size=arg.patch,
window=11,
radius=5)
if torch.cuda.is_available():
model.cuda()
torch.backends.cudnn.benchmark = True
data = DataLoader(obj_data,
batch_size=arg.batch,
shuffle=True,
num_workers=arg.workers,
pin_memory=True)
else:
data = DataLoader(obj_data, batch_size=arg.batch, shuffle=True)
# network parameters
epsilon = np.sqrt(1.0e-6)
p_window = 7
if clipping:
gadient_clip_Dnet = 1000.0
gadient_clip_Snet = 50.0
Dnet_parameters = [
x for name, x in model.named_parameters()
if 'dnet' in name.lower()
]
Snet_parameters = [
x for name, x in model.named_parameters()
if 'snet' in name.lower()
]
optimizer = optim.Adam(model.parameters(), lr=2e-4)
scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=arg.gamma)
print("Training model Benchmark now!")
for epoch in range(arg.epochs):
tic = time.time()
if clipping:
grad_D = 0.0
grad_S = 0.0
epoch_avg_loss = 0.0
mse_avg = 0.0
psnr_avg = 0.0
ssim_avg = 0.0
lr = optimizer.param_groups[0]['lr']
if lr < arg.learning:
print("reach min learning rate at epoch" + str(epoch))
model.train()
for i, batch_data in enumerate(data):
if torch.cuda.is_available():
y_batch, x_batch, sigma_arr = Variable(
batch_data[0]).cuda(), Variable(
batch_data[1]).cuda(), Variable(batch_data[2]).cuda()
else:
y_batch, x_batch, sigma_arr = batch_data[0], batch_data[
1], batch_data[2]
optimizer.zero_grad()
out_D, out_s = model(y_batch)
loss, loglikelihood, kl_z, kl_sigma = loss_func.get_loss(
x_batch, y_batch, sigma_arr, p_window,
out_D[:, :arg.channels, :, :], out_D[:, arg.channels:, :, :],
out_s[:, :arg.channels, :, :], out_s[:, arg.channels:, :, :],
epsilon)
loss.backward()
if clipping:
full_grad_D = nn.utils.clip_grad_norm_(Dnet_parameters,
gadient_clip_Dnet)
full_grad_S = nn.utils.clip_grad_norm_(Snet_parameters,
gadient_clip_Snet)
grad_D = (grad_D * (i / (i + 1)) + full_grad_D / (i + 1))
grad_S = (grad_S * (i / (i + 1)) + full_grad_S / (i + 1))
optimizer.step()
epoch_avg_loss += loss.detach().item()
predicted_image = y_batch - out_D[:, :arg.channels, :, :].detach(
).data
predicted_image = predicted_image.clamp(0, 1)
mse = calc_MSE(predicted_image, x_batch)
mse_avg += mse
psnr_avg += psnr(predicted_image * 255, x_batch * 255)
ssim_avg += calculate_ssim(img_as_ubyte(
predicted_image.permute(2, 3, 1, 0).cpu().numpy()),
img_as_ubyte(
x_batch.permute(2, 3, 1,
0).cpu().numpy()),
multichannel=True)
if i == 0:
print("First ForwardPAss\n Loss: {}, MSE: {}".format(
loss.detach().item(), mse))
if (i + 1) % 100 == 0:
print("{} - Loss: {}, MSE:{}, epoch:{}".format(
i + 1, loss.item(), mse, epoch + 1))
if i >= 5000:
break
if clipping:
gadient_clip_Dnet = min(gadient_clip_Dnet, grad_D)
gadient_clip_Dnet = min(gadient_clip_Dnet, grad_S)
print("----------------------------------------------------------")
print(
"Epoch: {}, Avg MSE:{}, Avg Epoch Loss:{}, Avg PSNR:{}, Avg SSIM : {}, LR:{}"
.format(epoch + 1, mse_avg / (i + 1), epoch_avg_loss / (i + 1),
psnr_avg / (i + 1), ssim_avg / (i + 1), lr))
logging.info("av loss: {}, epoch: {}".format(epoch_avg_loss / (i + 1),
epoch + 1))
# --------------- here comes the validation! ---------------
model.eval()
avg_psnr_validation = 0.0
avg_ssim_validation = 0.0
obj_data = gd.ValidationBenchmark(h5_file_=arg.val_data)
if torch.cuda.is_available():
model.cuda()
torch.backends.cudnn.benchmark = True
for idx in range(obj_data.__len__()):
noisy, image = obj_data.__getitem__(idx)
ch, ht, wt = noisy.shape
noisy = noisy.view(1, ch, ht, wt).cuda()
image = image.cuda()
model_out, _ = model(noisy)
noise = noisy - model_out[:, :ch, ].detach().data
clean_img_pred = noise.view(ch, ht, wt).permute(1, 2,
0).clamp(0, 1)
image = image.view(ch, ht, wt).permute(1, 2, 0)
avg_psnr_validation += psnr(image * 255, clean_img_pred * 255)
avg_ssim_validation += compare_ssim(
img_as_ubyte(image.cpu().numpy()),
img_as_ubyte(clean_img_pred.cpu().numpy()),
win_size=11,
data_range=255,
multichannel=True,
gaussian_weights=True)
print("average validation PSNR = ",
avg_psnr_validation / obj_data.__len__())
print("average validation SSIM = ",
avg_ssim_validation / obj_data.__len__())
# -------------- finish validation ---------------------------------
scheduler.step()
toc = time.time()
print('Time for this epoch: {:.2f}'.format(toc - tic))
if epoch % arguments.epoch_save == 0:
torch.save(
model.state_dict(),
os.path.join(arg.model_path,
"model_" + str(epoch) + "_epochs.pth"))
print("saved model as" + arg.model_path)
print("Finished Training...\n Saving model now.....\n")
torch.save(model.state_dict(),
os.path.join(arg.model_path, "final_model.pth"))
print("saved model as" + os.path.join(arg.model_path, "final_model.pth"))