-
Notifications
You must be signed in to change notification settings - Fork 0
/
model_unetskip.py
147 lines (129 loc) · 6.09 KB
/
model_unetskip.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
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
144
145
146
147
import torch
import torch.nn as nn
class ReconstructiveSubNetwork(nn.Module):
def __init__(self,in_channels=3, out_channels=3, base_width=128):
super(ReconstructiveSubNetwork, self).__init__()
self.encoder = EncoderReconstructive(in_channels, base_width)
self.decoder = DecoderReconstructive(base_width, out_channels=out_channels)
def forward(self, x):
b5,b4,b3,b2,b1 = self.encoder(x)
output = self.decoder(b5,b4,b3,b2,b1)
return output
class EncoderReconstructive(nn.Module):
def __init__(self, in_channels, base_width):
super(EncoderReconstructive, self).__init__()
self.block1 = nn.Sequential(
nn.Conv2d(in_channels,base_width, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width),
nn.ReLU(inplace=True),
nn.Conv2d(base_width, base_width, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width),
nn.ReLU(inplace=True))
self.mp1 = nn.Sequential(nn.MaxPool2d(2))
self.block2 = nn.Sequential(
nn.Conv2d(base_width,base_width*2, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*2),
nn.ReLU(inplace=True),
nn.Conv2d(base_width*2, base_width*2, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*2),
nn.ReLU(inplace=True))
self.mp2 = nn.Sequential(nn.MaxPool2d(2))
self.block3 = nn.Sequential(
nn.Conv2d(base_width*2,base_width*4, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*4),
nn.ReLU(inplace=True),
nn.Conv2d(base_width*4, base_width*4, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*4),
nn.ReLU(inplace=True))
self.mp3 = nn.Sequential(nn.MaxPool2d(2))
self.block4 = nn.Sequential(
nn.Conv2d(base_width*4,base_width*8, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*8),
nn.ReLU(inplace=True),
nn.Conv2d(base_width*8, base_width*8, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*8),
nn.ReLU(inplace=True))
self.mp4 = nn.Sequential(nn.MaxPool2d(2))
self.block5 = nn.Sequential(
nn.Conv2d(base_width*8,base_width*8, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*8),
nn.ReLU(inplace=True),
nn.Conv2d(base_width*8, base_width*8, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*8),
nn.ReLU(inplace=True))
def forward(self, x):
b1 = self.block1(x)
mp1 = self.mp1(b1)
b2 = self.block2(mp1)
mp2 = self.mp2(b2)
b3 = self.block3(mp2)
mp3 = self.mp3(b3)
b4 = self.block4(mp3)
mp4 = self.mp4(b4)
b5 = self.block5(mp4)
return b5,b4,b3,b2,b1
class DecoderReconstructive(nn.Module):
def __init__(self, base_width, out_channels=1):
super(DecoderReconstructive, self).__init__()
self.up1 = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(base_width * 8, base_width * 8, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width * 8),
nn.ReLU(inplace=True))
self.db1 = nn.Sequential(
nn.Conv2d(base_width*8, base_width*8, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*8),
nn.ReLU(inplace=True),
nn.Conv2d(base_width * 8, base_width * 4, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width * 4),
nn.ReLU(inplace=True)
)
self.up2 = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(base_width * 4, base_width * 4, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width * 4),
nn.ReLU(inplace=True))
self.db2 = nn.Sequential(
nn.Conv2d(base_width*4, base_width*4, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*4),
nn.ReLU(inplace=True),
nn.Conv2d(base_width * 4, base_width * 2, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width * 2),
nn.ReLU(inplace=True)
)
self.up3 = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(base_width * 2, base_width*2, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*2),
nn.ReLU(inplace=True))
# cat with base*1
self.db3 = nn.Sequential(
nn.Conv2d(base_width*2, base_width*2, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*2),
nn.ReLU(inplace=True),
nn.Conv2d(base_width*2, base_width*1, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width*1),
nn.ReLU(inplace=True)
)
self.up4 = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(base_width, base_width, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width),
nn.ReLU(inplace=True))
self.db4 = nn.Sequential(
nn.Conv2d(base_width*1, base_width, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width),
nn.ReLU(inplace=True),
nn.Conv2d(base_width, base_width, kernel_size=3, padding=1),
nn.BatchNorm2d(base_width),
nn.ReLU(inplace=True)
)
self.fin_out = nn.Sequential(nn.Conv2d(base_width, out_channels, kernel_size=3, padding=1))
#self.fin_out = nn.Conv2d(base_width, out_channels, kernel_size=3, padding=1)
def forward(self, b5,b4,b3,b2,b1):
up1 = self.up1(b5)
db1 = self.db1(up1+b4)
up2 = self.up2(db1)
db2 = self.db2(up2+b3)
up3 = self.up3(db2)
db3 = self.db3(up3+b2)
up4 = self.up4(db3)
db4 = self.db4(up4+b1)
out = self.fin_out(db4)
return out