-
Notifications
You must be signed in to change notification settings - Fork 0
/
Denoiser.py
109 lines (102 loc) · 4.08 KB
/
Denoiser.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from __future__ import division
import torch
import torch.nn as nn
from torch.autograd import Variable
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
import time
import numpy as np
from CNN import Net
from skimage.restoration import estimate_sigma
import cv2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def weights_init_kaiming(m):
class_name = m.__class__.__name__
if class_name.find('Linear') != -1:
nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif class_name.find('Conv2d') != -1:
nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif class_name.find('Norm') != -1:
m.weight.data.normal_(1.0, 0.02)
if m.bias is not None:
m.bias.data.zero_()
def Denoiser(V, dim, k_subspace, num_steps, mode = 'pre-train', model_path = './models/model.pkl'):
# dimensionality reduction
v = np.reshape(V, [dim[0], dim[1] * dim[2]])
u, _, _ = np.linalg.svd(v, full_matrices=False)
u = u[:, 0:k_subspace]
eigen_v = np.dot(u.transpose(), v)
# normalization
Min = np.amin(eigen_v, axis = 1)
Max = np.amax(eigen_v, axis = 1)
scale = Max - Min
scale = np.tile(scale, [dim[1] * dim[2], 1]).transpose()
Min = np.tile(Min, [dim[1] * dim[2], 1]).transpose()
eigen_v = (eigen_v - Min) / scale
eigen_V = np.reshape(eigen_v, [k_subspace, dim[1], dim[2]])
# cv2.imshow("Resized image", V[0,:,:])
# cv2.waitKey(0)
# cv2.destroyAllWindows()
V = np.expand_dims(eigen_V, axis = 0)
image = Variable(torch.from_numpy(V).float().to(device))
# initialise the network
net = Net(nLayers = 8, nChannels = k_subspace, multiply = 4).to(device)
if mode == 'pre-train' or mode == 'fine-tune':
# input of the network is image + noise
sigma_est = np.array(estimate_sigma(np.transpose(eigen_V, [1,2,0]), average_sigmas=False, multichannel=True))
# print('sigma_est: ', sigma_est)
sigma_est = Variable(torch.from_numpy(np.expand_dims(np.transpose(np.tile(sigma_est, [dim[1], dim[2], 1]), [2,0,1]), axis = 0)).float().to(device))
# network optimizer set up
optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)
if mode == 'pre-train':
net.apply(weights_init_kaiming)
else:
checkpoint = torch.load(model_path)
net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
loss = checkpoint['loss']
net.train()
# optimize
for step in range(num_steps):
# generate noise
noise = Variable(sigma_est * torch.randn(np.shape(V)).float().to(device))
# get the network output
output = net(image + noise)
# calculate the l2_loss over the masked output and take an optimizer step
optimizer.zero_grad()
loss = torch.mean(torch.abs(output - image))
loss.backward()
optimizer.step()
if (step == 0) or ((step + 1) % 1000 == 0):
print('At step {}, loss is {}'.format(step + 1, loss.data.cpu()))
# save model
torch.save({
'model_state_dict': net.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
}, model_path)
elif mode == 'test':
checkpoint = torch.load(model_path)
net.load_state_dict(checkpoint['model_state_dict'])
else:
raise SystemExit('error!')
# output image
net.eval()
output = net(image)
output = output.to('cpu')
output = output.detach().numpy()
output = np.squeeze(output)
# reconstruct data using denoising engin images
output = np.clip(output, 0, 1)
output = np.reshape(output, [k_subspace, dim[1]*dim[2]])
output = output * scale + Min
output = np.dot(u, output)
output = np.reshape(output, dim)
# clean up any mess we're leaving on the gpu
if torch.cuda.is_available():
torch.cuda.empty_cache()
return output