-
Notifications
You must be signed in to change notification settings - Fork 11
/
utils.py
95 lines (73 loc) · 2.6 KB
/
utils.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
import torch
import random
from math import log10
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm2d') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0.0)
# recommend
def initialize_weights(m):
if isinstance(m, nn.Conv2d):
# m.weight.data.normal_(0, 0.02)
# m.bias.data.zero_()
# nn.init.xavier_normal_(m.weight.data)
nn.init.kaiming_normal(m.weight.data, mode='fan_out')
# nn.init.xavier_normal_(m.bias.data)
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.02)
m.bias.data.zero_()
class AverageMeter():
""" Computes and stores the average and current value """
def __init__(self):
self.reset()
def reset(self):
""" Reset all statistics """
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
""" Update statistics """
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def to_psnr(J, gt):
mse = F.mse_loss(J, gt, reduction='none')
mse_split = torch.split(mse, 1, dim=0)
mse_list = [torch.mean(torch.squeeze(mse_split[ind])).item() for ind in range(len(mse_split))]
intensity_max = 1.0
psnr_list = [10.0 * log10(intensity_max / mse) for mse in mse_list]
return psnr_list
def create_emamodel(net, ema=True):
if ema:
for param in net.parameters():
param.detach_()
return net
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def compute_psnr_ssim(recoverd, clean):
assert recoverd.shape == clean.shape
recoverd = np.clip(recoverd.detach().cpu().numpy(), 0, 1)
clean = np.clip(clean.detach().cpu().numpy(), 0, 1)
recoverd = recoverd.transpose(0, 2, 3, 1)
clean = clean.transpose(0, 2, 3, 1)
psnr = 0
ssim = 0
for i in range(recoverd.shape[0]):
psnr += peak_signal_noise_ratio(clean[i], recoverd[i], data_range=1)
ssim += structural_similarity(clean[i], recoverd[i], data_range=1, multichannel=True)
return psnr / recoverd.shape[0], ssim / recoverd.shape[0], recoverd.shape[0]