ProbUNet-Tutorial/data.py at main · mhsotoudeh/ProbUNet-Tutorial · 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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import pickle

import numpy as np
import torchvision.transforms as transforms
from torch.utils.data import Dataset


class GRFDataset(Dataset):
    def __init__(self, inputs, truths, img_transform=None, truth_transform=None, metas=None):
        self.inputs = inputs
        self.truths = truths
        self.img_transform = img_transform
        self.truth_transform = truth_transform
        self.metas = metas

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, idx):
        image = self.inputs[idx]
        truth = self.truths[idx]

        if self.img_transform:
            image = self.img_transform(image)
        if self.truth_transform:
            truth = self.truth_transform(truth)

        if self.metas is None:
            return image, truth
        else:
            meta = self.metas[idx]
            return image, truth, meta


def prepare_data(datafile, size=None, shuffle=False, normalization='standard', eps=1e-6, metafile=None):
    with open(datafile, 'rb') as f:
        truths = np.load(f)
        inputs = np.load(f)

    if size is None:
        size = len(inputs)

    if shuffle is True:
        p = np.random.permutation(len(inputs))
        inputs, truths = inputs[p], truths[p]

    # Select required data
    inputs  = inputs[:size]
    truths = truths[:size]

    # Load Meta
    metas = None
    if metafile is not None:
        with open(metafile, 'rb') as f:
            metas = pickle.load(f)
        if shuffle is True:
            metas = np.array(metas)[p].tolist()


    # Define Normalization Transforms
    if normalization == 'standard':
        # Calculate mean and std for standardization
        input_mean, input_std = np.mean(inputs), np.std(inputs)
        truth_mean, truth_std = np.mean(truths), np.std(truths)


        # Define Transforms
        input_trans = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=input_mean, std=input_std)
        ])

        truth_trans = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=truth_mean, std=truth_std)
        ])


        # Define Inverse Transforms
        inv_input_trans = transforms.Compose([
            transforms.Normalize(mean=0.0, std=1.0/input_std),
            transforms.Normalize(mean=-input_mean, std=1.0)
        ])

        inv_truth_trans = transforms.Compose([
            transforms.Normalize(mean=0.0, std=1.0/truth_std),
            transforms.Normalize(mean=-truth_mean, std=1.0)
        ])


    elif normalization == 'normal':
        # Calculate min and max for normalization
        input_min, input_max = np.mean(inputs), np.std(inputs)
        truth_min, truth_max = np.mean(truths), np.std(truths)


        # Define Transforms
        input_trans = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=input_mean, std=input_std)
        ])

        truth_trans = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=truth_mean, std=truth_std)
        ])


        # Define Inverse Transforms
        inv_input_trans = transforms.Compose([
            transforms.Normalize(mean=0.0, std=1.0/input_std),
            transforms.Normalize(mean=-input_mean, std=1.0)
        ])

        inv_truth_trans = transforms.Compose([
            transforms.Normalize(mean=0.0, std=1.0/truth_std),
            transforms.Normalize(mean=-truth_mean, std=1.0)
        ])


    elif normalization is None:
        # Define Transforms
        input_trans = transforms.ToTensor()
        truth_trans = transforms.ToTensor()

        # Define Inverse Transforms
        inv_input_trans = transforms.Normalize(mean=0.0, std=1.0)
        inv_truth_trans = transforms.Normalize(mean=0.0, std=1.0)


    transdict = {
        'input_transform': input_trans,
        'truth_transform': truth_trans,
        'inv_input_transform': inv_input_trans,
        'inv_truth_transform': inv_truth_trans
    }


    # Create Datasets
    train_data = GRFDataset(inputs, truths, img_transform=input_trans, truth_transform=truth_trans, metas=metas)

    return train_data, transdict