if opt.arch == 'DnCNN-S':
model = DnCNN(num_layers=17)
elif opt.arch == 'DnCNN-B':
model = DnCNN(num_layers=20)
elif opt.arch == 'DnCNN-3':
model = DnCNN(num_layers=20)
state_dict = model.state_dict()
for n, p in torch.load(opt.weights_path,
map_location=lambda storage, loc: storage).items():
if n in state_dict.keys():
state_dict[n].copy_(p)
else:
raise KeyError(n)
model = model.to(device)
model.eval()
filename = os.path.basename(opt.image_path).split('.')[0]
descriptions = ''
input = pil_image.open(opt.image_path).convert('RGB')
if opt.gaussian_noise_level is not None:
noise = np.random.normal(0.0, opt.gaussian_noise_level,
(input.height, input.width, 3)).astype(
np.float32)
input = np.array(input).astype(np.float32) + noise
descriptions += '_noise_l{}'.format(opt.gaussian_noise_level)
pil_image.fromarray(input.clip(0.0, 255.0).astype(np.uint8)).save(
os.path.join(opt.outputs_dir,
def train_model(config):
# Define hyper-parameters.
depth = int(config["DnCNN"]["depth"])
n_channels = int(config["DnCNN"]["n_channels"])
img_channel = int(config["DnCNN"]["img_channel"])
kernel_size = int(config["DnCNN"]["kernel_size"])
use_bnorm = config.getboolean("DnCNN", "use_bnorm")
epochs = int(config["DnCNN"]["epoch"])
batch_size = int(config["DnCNN"]["batch_size"])
train_data_dir = config["DnCNN"]["train_data_dir"]
test_data_dir = config["DnCNN"]["test_data_dir"]
eta_min = float(config["DnCNN"]["eta_min"])
eta_max = float(config["DnCNN"]["eta_max"])
dose = float(config["DnCNN"]["dose"])
model_save_dir = config["DnCNN"]["model_save_dir"]
# Save logs to txt file.
log_dir = config["DnCNN"]["log_dir"]
log_dir = Path(log_dir) / "dose{}".format(str(int(dose * 100)))
log_file = log_dir / "train_result.txt"
# Define device.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Initiate a DnCNN instance.
# Load the model to device and set the model to training.
model = DnCNN(depth=depth, n_channels=n_channels,
img_channel=img_channel,
use_bnorm=use_bnorm,
kernel_size=kernel_size)
model = model.to(device)
model.train()
# Define loss criterion and optimizer
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2)
criterion = LossFunc(reduction="mean")
# Get a validation test set and corrupt with noise for validation performance.
# For every epoch, use this pre-determined noisy images.
test_file_list = glob.glob(test_data_dir + "/*.png")
xs_test = []
# Can't directly convert the xs_test from list to ndarray because some images are 512*512
# while the rest are 256*256.
for i in range(len(test_file_list)):
img = cv2.imread(test_file_list[i], 0)
img = np.array(img, dtype="float32") / 255.0
img = np.expand_dims(img, axis=0)
img_noisy, _ = nm(img, eta_min, eta_max, dose, t=100)
xs_test.append((img_noisy, img))
# Train the model.
loss_store = []
epoch_loss_store = []
psnr_store = []
ssim_store = []
psnr_tr_store = []
ssim_tr_store = []
loss_mse = torch.nn.MSELoss()
dtype = torch.cuda.FloatTensor
# load vgg network
vgg = Vgg16().type(dtype)
for epoch in range(epochs):
# For each epoch, generate clean augmented patches from the training directory.
# Convert the data from uint8 to float32 then scale them to make it in [0, 1].
# Then make the patches to be of shape [N, C, H, W],
# where N is the batch size, C is the number of color channels.
# H and W are height and width of image patches.
xs = dg.datagenerator(data_dir=train_data_dir)
xs = xs.astype("float32") / 255.0
xs = torch.from_numpy(xs.transpose((0, 3, 1, 2)))
train_set = dg.DenoisingDatatset(xs, eta_min, eta_max, dose)
train_loader = DataLoader(dataset=train_set, num_workers=4,
drop_last=True, batch_size=batch_size,
shuffle=True) # TODO: if drop_last=True, the dropping in the
# TODO: data_generator is not necessary?
# train_loader_test = next(iter(train_loader))
t_start = timer()
epoch_loss = 0
for idx, data in enumerate(train_loader):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
img_batch_read = len(inputs)
optimizer.zero_grad()
outputs = model(inputs)
# We can use labels for both style and content image
# style image
# style_transform = transforms.Compose([
# normalize_tensor_transform() # normalize with ImageNet values
# ])
# labels_t = style_transform(labels)
labels_t = labels.repeat(1, 3, 1, 1)
outputs_t = outputs.repeat(1, 3, 1, 1)
y_c_features = vgg(labels_t)
style_gram = [gram(fmap) for fmap in y_c_features]
y_hat_features = vgg(outputs_t)
y_hat_gram = [gram(fmap) for fmap in y_hat_features]
# calculate style loss
style_loss = 0.0
for j in range(4):
style_loss += loss_mse(y_hat_gram[j], style_gram[j][:img_batch_read])
style_loss = STYLE_WEIGHT*style_loss
aggregate_style_loss = style_loss
# calculate content loss (h_relu_2_2)
recon = y_c_features[1]
recon_hat = y_hat_features[1]
content_loss = CONTENT_WEIGHT*loss_mse(recon_hat, recon)
aggregate_content_loss = content_loss
loss = aggregate_content_loss + aggregate_style_loss
# loss = criterion(outputs, labels)
loss_store.append(loss.item())
epoch_loss += loss.item()
loss.backward()
optimizer.step()
if idx % 100 == 0:
print("Epoch [{} / {}], step [{} / {}], loss = {:.5f}, lr = {:.6f}, elapsed time = {:.2f}s".format(
epoch + 1, epochs, idx, len(train_loader), loss.item(), *scheduler.get_last_lr(), timer()-t_start))
epoch_loss_store.append(epoch_loss / len(train_loader))
# At each epoch validate the result.
model = model.eval()
# # Firstly validate on training sets. This takes a long time so I commented.
# tr_psnr = []
# tr_ssim = []
# # t_start = timer()
# with torch.no_grad():
# for idx, train_data in enumerate(train_loader):
# inputs, labels = train_data
# # print(inputs.shape)
# # inputs = np.expand_dims(inputs, axis=0)
# # inputs = torch.from_numpy(inputs).to(device)
# inputs = inputs.to(device)
# labels = labels.squeeze().numpy()
#
# outputs = model(inputs)
# outputs = outputs.squeeze().cpu().detach().numpy()
#
# tr_psnr.append(peak_signal_noise_ratio(labels, outputs))
# tr_ssim.append(structural_similarity(outputs, labels))
# psnr_tr_store.append(sum(tr_psnr) / len(tr_psnr))
# ssim_tr_store.append(sum(tr_ssim) / len(tr_ssim))
# # print("Elapsed time = {}".format(timer() - t_start))
#
# print("Validation on train set: epoch [{} / {}], aver PSNR = {:.2f}, aver SSIM = {:.4f}".format(
# epoch + 1, epochs, psnr_tr_store[-1], ssim_tr_store[-1]))
# Validate on test set
val_psnr = []
val_ssim = []
with torch.no_grad():
for idx, test_data in enumerate(xs_test):
inputs, labels = test_data
inputs = np.expand_dims(inputs, axis=0)
inputs = torch.from_numpy(inputs).to(device)
labels = labels.squeeze()
outputs = model(inputs)
outputs = outputs.squeeze().cpu().detach().numpy()
val_psnr.append(peak_signal_noise_ratio(labels, outputs))
val_ssim.append(structural_similarity(outputs, labels))
psnr_store.append(sum(val_psnr) / len(val_psnr))
ssim_store.append(sum(val_ssim) / len(val_ssim))
print("Validation on test set: epoch [{} / {}], aver PSNR = {:.2f}, aver SSIM = {:.4f}".format(
epoch + 1, epochs, psnr_store[-1], ssim_store[-1]))
# Set model to train mode again.
model = model.train()
scheduler.step()
# Save model
save_model(model, model_save_dir, epoch, dose * 100)
# Save the loss and validation PSNR, SSIM.
if not log_dir.exists():
Path.mkdir(log_dir)
with open(log_file, "a+") as fh:
# fh.write("{} Epoch [{} / {}], loss = {:.6f}, train PSNR = {:.2f}dB, train SSIM = {:.4f}, "
# "validation PSNR = {:.2f}dB, validation SSIM = {:.4f}".format(
# datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"),
# epoch + 1, epochs, epoch_loss_store[-1],
# psnr_tr_store[-1], ssim_tr_store[-1],
# psnr_store[-1], ssim_store[-1]))
fh.write("{} Epoch [{} / {}], loss = {:.6f}, "
"validation PSNR = {:.2f}dB, validation SSIM = {:.4f}\n".format(
datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"),
epoch + 1, epochs, epoch_loss_store[-1],
psnr_store[-1], ssim_store[-1]))
# np.savetxt(log_file, np.hstack((epoch + 1, epoch_loss_store[-1], psnr_store[-1], ssim_store[-1])), fmt="%.6f", delimiter=", ")
fig, ax = plt.subplots()
ax.plot(loss_store[-len(train_loader):])
ax.set_title("Last 1862 losses")
ax.set_xlabel("iteration")
fig.show()