GenerativeCardiacMotion/model.py at main · shaharzuler/GenerativeCardiacMotion · 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
import torch
import torch.nn as nn


def conv(in_planes, out_planes, kernel_size=3, stride=1, dilation=1, isReLU=True,
         conv_module=nn.Conv3d):
    if isReLU:
        return nn.Sequential(
            conv_module(in_planes, out_planes, kernel_size=kernel_size,
                        stride=stride, dilation=dilation,
                        padding=((kernel_size - 1) * dilation) // 2, bias=True),
            nn.LeakyReLU(0.1, inplace=True)
        )
    else:
        return nn.Sequential(
            nn.Conv3d(in_planes, out_planes, kernel_size=kernel_size,
                      stride=stride, dilation=dilation,
                      padding=((kernel_size - 1) * dilation) // 2, bias=True)
        )


class FeatureExtractor(nn.Module):
    def __init__(self, num_chs):
        super(FeatureExtractor, self).__init__()
        self.num_chs = num_chs
        self.convs = nn.ModuleList()

        for l, (ch_in, ch_out) in enumerate(zip(num_chs[:-1], num_chs[1:])):
            layer = nn.Sequential(
                conv(ch_in, ch_out, stride=2),
                conv(ch_out, ch_out)
            )
            self.convs.append(layer)

    def forward(self, x):
        feature_pyramid = []
        for conv in self.convs:
            x = conv(x)
            feature_pyramid.append(x)
        return feature_pyramid[::-1]


class Encoder(nn.Module):
    def __init__(self, num_chs, max_var=6):
        super(Encoder, self).__init__()
        self.num_chs = num_chs
        self.max_var = max_var
        self.epsilon = 1e-8
        self.convs = nn.ModuleList()
        kernel_size = 3
        stride = 2
        for l, (ch_in, ch_out) in enumerate(zip(num_chs[:-1], num_chs[1:])):
            if l == 0:
                ch_in_expanded = ch_in + 1
            else:
                ch_in_expanded = ch_in * 2
            if l == 5:
                kernel_size = 1
                stride = 1
            layer = nn.Sequential(
                conv(ch_in_expanded, ch_out, stride=stride, kernel_size=kernel_size),
                conv(ch_out, ch_out, kernel_size=kernel_size)
            )
            self.convs.append(layer)
        self.mean = nn.Conv3d(num_chs[-1], num_chs[-1], kernel_size=1, stride=1, dilation=1, padding=0 // 2, bias=True)
        self.var = nn.Conv3d(num_chs[-1], num_chs[-1], kernel_size=1, stride=1, dilation=1, padding=0 // 2, bias=True)
        self.kl = 0

    def forward(self, flow, condition_pyramid):
        x = flow
        for conv, condition_features in zip(self.convs, condition_pyramid[::-1]):  # check
            x = torch.concat((x, condition_features), dim=1)
            x = conv(x)
        mean = self.mean(x)
        std = torch.exp(0.5 * torch.clamp(self.var(x), max=self.max_var))
        eps = torch.randn_like(std)
        z = mean + eps * std
        self.kl = (std ** 2 + mean ** 2 - torch.log(std + self.epsilon) - 1 / 2).mean()

        return z


class Decoder(nn.Module):
    def __init__(self, num_chs, max_flow_hat_abs_val):
        super(Decoder, self).__init__()
        self.num_chs = num_chs
        self.max_f_hat_abs_val = abs(max_flow_hat_abs_val)
        self.convs = nn.ModuleList()
        kernel_size = 3
        stride = 2
        for l, (ch_in, ch_out) in enumerate(zip(num_chs[:-1], num_chs[1:])):
            if l == 5:
                ch_in_expanded = ch_in + 1
                kernel_size = 1
                stride = 1
            else:
                ch_in_expanded = ch_in * 2
            layer = nn.Sequential(
                conv(ch_in_expanded, ch_out, stride=stride, kernel_size=kernel_size, conv_module=nn.ConvTranspose3d),
                conv(ch_out, ch_out, kernel_size=kernel_size, conv_module=nn.ConvTranspose3d, isReLU=False)
            )
            if l == 0:
                layer[0][0].output_padding = (1, 1, 0)
            if l == 1:
                layer[0][0].output_padding = (0, 0, 1)
            if l == 2:
                layer[0][0].output_padding = (1, 1, 0)
            if l == 3:
                layer[0][0].output_padding = (0,0,1)
            if l == 4:
                layer[0][0].output_padding = (1,1,1)#(0, 0, 1)
            self.convs.append(layer)

    def forward(self, z, condition_pyramid):
        x = z
        for conv, condition_features in zip(self.convs, condition_pyramid):  # check
            x = torch.concat((x, condition_features), dim=1)
            x = conv(x)
        x = torch.clip(x, min=-self.max_f_hat_abs_val, max=self.max_f_hat_abs_val)
        return x


class CVAE(torch.nn.Module):
    def __init__(self, num_chs, max_flow_hat_abs_val=50):
        super().__init__()
        self.encoder = Encoder(num_chs=num_chs)
        self.decoder = Decoder(num_chs=num_chs[:-1][::-1] + [3], max_flow_hat_abs_val=max_flow_hat_abs_val)

    def forward(self, flow, conditions_pyramid):
        z = self.encoder(flow.float(), conditions_pyramid)
        flow_hat = self.decoder(z, conditions_pyramid)
        return flow_hat, self.encoder.kl

    def generate(self, conditions_pyramid, device, latent_shape=(1,128,3,3,3), output_z=False, provide_z=False):
        if type(provide_z) == bool and not provide_z:
            z = torch.randn(latent_shape).to(device)
        else:
            z = provide_z
        flow_hat = self.decoder(z, conditions_pyramid)
        if not output_z:
            return flow_hat
        return flow_hat, z.flatten()