ModifiedDnCNN/utils.py at master · BiancaBing/ModifiedDnCNN · GitHub

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
import torch
import torch.nn as nn
import torch.nn.init as init
import numpy as np
from torch.nn.modules.loss import _Loss
import math

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def initialize_weights(m):
    if isinstance(m, nn.Conv2d):
        init.xavier_uniform_(m.weight, gain=1)
        # init.normal_(m.weight, std=0.001)
        if m.bias is not None:
            init.constant_(m.bias, 0)

def params_count(model):
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    return params

class mean_squared_error(_Loss):
    def __init__(self):
        super(mean_squared_error, self).__init__()

    def forward(self, input, target):
        # _assert_no_grad(target)
        return torch.nn.functional.mse_loss(input, target)

def get_full_name_list(list_name, stage_flag):
    full_image_list = []
    with open(list_name) as f:
        image_list = f.readlines()  # 00001/0001　-> 0001
    if stage_flag == 'train':
        str_print = 'len(train_image_list):'
    elif stage_flag == 'test':
        str_print = 'len(test_image_list):'
    print(str_print + str(len(image_list)))
    for i in range(0, int(len(image_list))):
        if stage_flag == 'train':
            if i % 2 == 0:
                til = image_list[i].rstrip()
                for j in range(1, 8):
                    til_png = til + '/im' + str(j) + '.png'
                    full_image_list.append(til_png)
        elif stage_flag == 'test':
            til = image_list[i].rstrip()
            for j in range(1, 8):
                til_png = til + '/im' + str(j) + '.png'
                full_image_list.append(til_png)
    if stage_flag == 'train':
        str_print = 'len(full_train_image_list):'
    elif stage_flag == 'test':
        str_print = 'len(full_test_image_list):'
    print(str_print + str(len(full_image_list)))
    return full_image_list

def calculate_psnr(input, prediction, label):
    psnrs = []
    mse_pre = math.sqrt(torch.mean((prediction - label) ** 2.0))
    psnr_pre = 20 * math.log10(1.0/mse_pre)
    mse_input = math.sqrt(torch.mean((input - label) ** 2.0))
    psnr_input = 20 * math.log10(1.0/mse_input)
    psnrs.append(psnr_pre)
    psnrs.append(psnr_input)
    return psnrs

def calculate_psnr_single(prediction, label):
    mse_pre = math.sqrt(torch.mean((prediction - label) ** 2.0))
    psnr_pre = 20 * math.log10(1.0/mse_pre)
    return psnr_pre