Skip to content

Commit eace23c

Browse files
Tensor46Tensor46
authored
Merge pull request #39 from Tensor46/develop
+ added anatomy net for segmentation
2 parents 7d3d63a + 099432a commit eace23c

File tree

6 files changed

+397
-185
lines changed

6 files changed

+397
-185
lines changed

‎core/NeuralArchitectures/__init__.py

+2-4
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
"ResidualNet", "InceptionV4", "DenseNet",
66
"LinearVAE", "ConvolutionalVAE", "PGGAN",
77
"ContextNet", "FeatureNet", "FeatureCapNet",
8-
"PointNet", "UNet", "UNetPatch", "UNetMini",
8+
"PointNet", "UNet", "ANet",
99
"NeuralDecisionForest", "Models"]
1010

1111
from .capsulenet import CapsuleNet
@@ -20,15 +20,13 @@
2020
from .convolutionalvae import ConvolutionalVAE
2121
from .pggan import PGGAN
2222
from .contextnet import ContextNet
23-
from .featurenet import FeatureNet, FeatureCapNet
2423
from .pointnet import PointNet
25-
from .unet import UNet, UNetMini
24+
from .unet import UNet, ANet
2625
from .trees import NeuralDecisionForest
2726

2827
del trees
2928
del unet
3029
del pointnet
31-
del featurenet
3230
del contextnet
3331
del pggan
3432
del capsulenet

‎core/NeuralArchitectures/contextnet.py

+111-36
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,126 @@
1-
2-
""" TensorMONK's :: NeuralArchitectures """
3-
4-
import torch
1+
""" TensorMONK's :: NeuralArchitectures """
52
import torch.nn as nn
6-
import numpy as np
7-
from ..NeuralLayers import *
8-
#==============================================================================#
3+
from ..NeuralLayers import Convolution
4+
from ..NeuralLayers import ContextNet_Bottleneck
95

106

117
class ContextNet(nn.Module):
12-
"""
13-
Implemented https://arxiv.org/pdf/1805.04554.pdf
8+
r"""ContextNet: Exploring Context and Detail for Semantic Seg in Real-time
9+
Implemented from https://arxiv.org/pdf/1805.04554.pdf
10+
11+
Naming Convention for layers in the module:
12+
cnv::Convolution:: regular convolution block
13+
bn::bottleneck:: bottlenect residual block
14+
dw::Convilution:: depth-wise seperable convolution block
15+
dn:: deep netwrk for Context
16+
sn:: shallow network for Context
17+
Args:
18+
tensor_size: shape of tensor in BCHW
1419
"""
1520
def __init__(self, tensor_size=(1, 3, 1024, 2048), *args, **kwargs):
1621
super(ContextNet, self).__init__()
17-
Strides = [2, 1, 1]
18-
bottleneck = ContextNet_Bottleneck
19-
normalization = "batch"
20-
self.DeepNET = nn.Sequential()
22+
normalization, strides = "batch", [2, 1, 1]
23+
bottleneck = ContextNet_Bottleneck
24+
25+
self.DeepNET = nn.Sequential()
2126
self.ShallowNET = nn.Sequential()
22-
self.DeepNET.add_module("AVGPL", nn.AvgPool2d((5,5), (4,4), 2)) # 1, 1, 256, 512
23-
self.DeepNET.add_module("DN_CNV1", Convolution(tensor_size, 3, 32, 2, True, "relu", 0., normalization, False, 1)) # 1, 1, 128, 256
24-
self.DeepNET.add_module("DN_BN10", bottleneck(self.DeepNET[-1].tensor_size, 3, 32, 1, expansion=1))
25-
self.DeepNET.add_module("DN_BN20", bottleneck(self.DeepNET[-1].tensor_size, 3, 32, 1, expansion=6)) # 1, 1, 128, 256
26-
for i in range(3): self.DeepNET.add_module("DN_BN3"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 48, Strides[i], expansion=6)) # 1, 1, 64, 128
27-
for i in range(3): self.DeepNET.add_module("DN_BN4"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 64, Strides[i], expansion=6)) # 1, 1, 32, 64
28-
for i in range(2): self.DeepNET.add_module("DN_BN5"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 96, 1, expansion=6))
29-
for i in range(2): self.DeepNET.add_module("DN_BN6"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 128, 1, expansion=6)) # 1, 1, 32, 64
30-
self.DeepNET.add_module("DN_CNV2", Convolution(self.DeepNET[-1].tensor_size, 3, 128, 1, True, "relu", 0., normalization, False, 1)) # 1, 1, 32, 64
31-
self.DeepNET.add_module("UPSMPLE", nn.Upsample(scale_factor = 4, mode = 'bilinear')) # 1, 1, 128, 256
32-
_tensor_size = (1, 128, self.DeepNET[-2].tensor_size[2]*4, self.DeepNET[-2].tensor_size[3]*4)
33-
self.DeepNET.add_module("DN_DW11", Convolution(_tensor_size, 3, _tensor_size[1], 1, True, "relu", 0., None, False, groups =_tensor_size[1], dilation = 4))
34-
self.DeepNET.add_module("DN_DW12", Convolution(self.DeepNET[-1].tensor_size, 1, 128, 1, True, "relu", 0., normalization, False, 1))
35-
self.DeepNET.add_module("DN_CNV3", Convolution(self.DeepNET[-1].tensor_size, 1, 128, 1, True, "relu", 0., normalization, False, 1)) # 128, 256
3627

28+
# 1, 1, 256, 512
29+
self.DeepNET.add_module("avgpl", nn.AvgPool2d((5, 5), (4, 4), 2))
30+
# 1, 1, 128, 256
31+
self.DeepNET.add_module("dn_cnv1", Convolution(tensor_size, 3, 32, 2,
32+
True, "relu", 0., normalization, False, 1))
33+
self.DeepNET.add_module("dn_bn10",
34+
bottleneck(self.DeepNET[-1].tensor_size, 3, 32,
35+
1, expansion=1))
36+
self.DeepNET.add_module("dn_bn20",
37+
bottleneck(self.DeepNET[-1].tensor_size, 3, 32,
38+
1, expansion=6)) # 1, 1, 128, 256
39+
for i in range(3):
40+
self.DeepNET.add_module("dn_bn3"+str(i),
41+
bottleneck(self.DeepNET[-1].tensor_size, 3,
42+
48, strides[i], expansion=6))
43+
# 1, 1, 64, 128
44+
for i in range(3):
45+
self.DeepNET.add_module("dn_bn4"+str(i),
46+
bottleneck(self.DeepNET[-1].tensor_size, 3,
47+
64, strides[i], expansion=6))
48+
# 1, 1, 32, 64
49+
for i in range(2):
50+
self.DeepNET.add_module("dn_bn5"+str(i),
51+
bottleneck(self.DeepNET[-1].tensor_size, 3,
52+
96, 1, expansion=6))
53+
for i in range(2):
54+
self.DeepNET.add_module("DN_BN6"+str(i),
55+
bottleneck(self.DeepNET[-1].tensor_size, 3,
56+
128, 1, expansion=6))
57+
# 1, 1, 32, 64
58+
self.DeepNET.add_module("dn_cnv2",
59+
Convolution(self.DeepNET[-1].tensor_size, 3,
60+
128, 1, True, "relu", 0.,
61+
normalization, False, 1))
62+
# 1, 1, 32, 64
63+
self.DeepNET.add_module("upsample", nn.Upsample(scale_factor=4,
64+
mode='bilinear'))
65+
# 1, 1, 128, 256
66+
_tensor_size = (1, 128, self.DeepNET[-2].tensor_size[2]*4,
67+
self.DeepNET[-2].tensor_size[3]*4)
68+
self.DeepNET.add_module("dn_dw11",
69+
Convolution(_tensor_size, 3, _tensor_size[1],
70+
1, True, "relu", 0., None, False,
71+
groups=_tensor_size[1],
72+
dilation=4))
73+
self.DeepNET.add_module("dn_dw12",
74+
Convolution(self.DeepNET[-1].tensor_size, 1,
75+
128, 1, True, "relu", 0.,
76+
normalization, False, 1))
77+
self.DeepNET.add_module("dn_cnv3",
78+
Convolution(self.DeepNET[-1].tensor_size, 1,
79+
128, 1, True, "relu", 0.,
80+
normalization, False, 1))
81+
# 128, 256
3782
activation, pre_nm, groups = "relu", False, 1
38-
self.ShallowNET.add_module("SM_CNV1", Convolution(tensor_size, 3, 32, 2, True, "relu", 0.,True, False, 1)) # 512 x 1024
39-
self.ShallowNET.add_module("SM_DW11", Convolution(self.ShallowNET[-1].tensor_size, 3, 32, 2, True, activation, 0., None, pre_nm, groups = tensor_size[1])) # 256, 512
40-
self.ShallowNET.add_module("SM_DW12", Convolution(self.ShallowNET[-1].tensor_size, 1, 64, 1,True, activation, 0., normalization, pre_nm, groups))
41-
self.ShallowNET.add_module("SM_DW21", Convolution(self.ShallowNET[-1].tensor_size, 3, 64, 2, True, activation, 0., None, pre_nm, groups = tensor_size[1])) # 128, 256
42-
self.ShallowNET.add_module("SM_DW22", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups))
43-
self.ShallowNET.add_module("SM_DW31", Convolution(self.ShallowNET[-1].tensor_size, 3, 128, 1, True, activation, 0., None, pre_nm, groups = tensor_size[1]))
44-
self.ShallowNET.add_module("SM_DW32", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups))
45-
self.ShallowNET.add_module("SM_CNV2", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups)) # 128, 256
83+
self.ShallowNET.add_module("sm_cnv1",
84+
Convolution(tensor_size, 3, 32, 2, True,
85+
"relu", 0., True, False, 1))
86+
# 512 x 1024
87+
self.ShallowNET.add_module("sm_dw11",
88+
Convolution(self.ShallowNET[-1].tensor_size,
89+
3, 32, 2, True, activation, 0.,
90+
None, pre_nm,
91+
groups=tensor_size[1]))
92+
# 256, 512
93+
self.ShallowNET.add_module("sm_dw12",
94+
Convolution(self.ShallowNET[-1].tensor_size,
95+
1, 64, 1, True, activation, 0.,
96+
normalization, pre_nm, groups))
97+
self.ShallowNET.add_module("sm_dw21",
98+
Convolution(self.ShallowNET[-1].tensor_size,
99+
3, 64, 2, True, activation, 0.,
100+
None, pre_nm,
101+
groups=tensor_size[1]))
102+
self.ShallowNET.add_module("sm_dw22",
103+
Convolution(self.ShallowNET[-1].tensor_size,
104+
1, 128, 1, True, activation, 0.,
105+
normalization, pre_nm, groups))
106+
self.ShallowNET.add_module("sm_dw31",
107+
Convolution(self.ShallowNET[-1].tensor_size,
108+
3, 128, 1, True, activation, 0.,
109+
None, pre_nm,
110+
groups=tensor_size[1]))
111+
self.ShallowNET.add_module("sm_dw32",
112+
Convolution(self.ShallowNET[-1].tensor_size,
113+
1, 128, 1, True, activation, 0.,
114+
normalization, pre_nm, groups))
115+
self.ShallowNET.add_module("sm_cnv2",
116+
Convolution(self.ShallowNET[-1].tensor_size,
117+
1, 128, 1, True, activation, 0.,
118+
normalization, pre_nm, groups))
119+
# 128, 256
46120
self.tensor_size = self.ShallowNET[-1].tensor_size
47121

48-
self.FuseNET = Convolution(self.ShallowNET[-1].tensor_size, 1, self.ShallowNET[-1].tensor_size[1], 1, True)
122+
self.FuseNET = Convolution(self.ShallowNET[-1].tensor_size, 1,
123+
self.ShallowNET[-1].tensor_size[1], 1, True)
49124

50125
def forward(self, tensor):
51126
return self.FuseNET(self.DeepNET(tensor)+self.ShallowNET(tensor))

‎core/NeuralArchitectures/pointnet.py

+41-23
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,58 @@
1+
""" TensorMONK's :: NeuralArchitectures """
2+
import torch.nn as nn
3+
from core.NeuralLayers import Convolution
4+
# =========================================================================== #
15

2-
""" TensorMONK's :: NeuralArchitectures """
36

4-
import torch
5-
import torch.nn as nn
6-
import numpy as np
7-
from ..NeuralLayers import *
8-
#==============================================================================#
97
class PointNet(nn.Module):
8+
r"""
9+
Implemented from paper: Learning Discriminative and Transformation
10+
Covariant Local Feature Detectors
11+
Args:
12+
tensor_size: shape of tensor in BCHW
13+
(None/any integer >0, channels, height, width)
14+
out_channels: depth of output feature channels.
15+
activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish
16+
normalization: None/batch/group/instance/layer/pixelwise
1017
"""
11-
Implemented http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhang_Learning_Discriminative_and_CVPR_2017_paper.pdf
12-
"""
13-
def __init__(self, tensor_size=(1, 1, 32, 32), out_channels=2, *args, **kwargs):
18+
def __init__(self, tensor_size=(1, 1, 32, 32), out_channels=2,
19+
*args, **kwargs):
1420
super(PointNet, self).__init__()
1521
normalization = "batch"
16-
self.PointNET = nn.Sequential()
17-
self.PointNET.add_module("CONV1", Convolution(tensor_size, 5, 32, 1, False, "relu", 0., None))
22+
activation = "relu"
23+
self.PointNET = nn.Sequential()
24+
self.PointNET.add_module("CONV1",
25+
Convolution(tensor_size, 5, 32, 1, False,
26+
activation, 0., None))
1827
self.PointNET.add_module("POOL1", nn.MaxPool2d(2))
1928
_tensor_size = self.PointNET[-2].tensor_size
20-
_tensor_size = (_tensor_size[0], _tensor_size[1], _tensor_size[2]//2, _tensor_size[3]//2)
21-
#print(_tensor_size)
22-
self.PointNET.add_module("CONV2", Convolution(_tensor_size, 3, 128,1, False, "relu", 0.,normalization))
29+
_tensor_size = (_tensor_size[0], _tensor_size[1],
30+
_tensor_size[2]//2, _tensor_size[3]//2)
31+
self.PointNET.add_module("CONV2",
32+
Convolution(_tensor_size, 3, 128, 1, False,
33+
activation, 0., normalization))
2334
self.PointNET.add_module("POOL2", nn.MaxPool2d(2))
2435
_tensor_size = self.PointNET[-2].tensor_size
25-
_tensor_size = (_tensor_size[0], _tensor_size[1], _tensor_size[2]//2, _tensor_size[3]//2)
26-
#print(_tensor_size)
27-
self.PointNET.add_module("CONV3", Convolution(_tensor_size, 3, 128, 1, False, "relu", 0.,normalization))
28-
#print(self.PointNET[-1].tensor_size)
29-
self.PointNET.add_module("CONV4", Convolution(self.PointNET[-1].tensor_size, 3, 256, 1, False, "relu", 0.,normalization))
30-
#print(self.PointNET[-1].tensor_size)
31-
self.PointNET.add_module("CONV5", Convolution(self.PointNET[-1].tensor_size, 2, out_channels, 1, False, "relu", 0.,None))
32-
print(self.PointNET[-1].tensor_size)
36+
_tensor_size = (_tensor_size[0], _tensor_size[1],
37+
_tensor_size[2]//2, _tensor_size[3]//2)
38+
self.PointNET.add_module("CONV3",
39+
Convolution(_tensor_size, 3, 128, 1, False,
40+
activation, 0., normalization))
41+
self.PointNET.add_module("CONV4",
42+
Convolution(self.PointNET[-1].tensor_size, 3,
43+
256, 1, False, activation, 0.,
44+
normalization))
45+
self.PointNET.add_module("CONV5",
46+
Convolution(self.PointNET[-1].tensor_size, 2,
47+
out_channels, 1, False,
48+
activation, 0., None))
3349
self.tensor_size = self.PointNET[-1].tensor_size
50+
3451
def forward(self, tensor):
3552
return self.PointNET(tensor).squeeze(2).squeeze(2)
3653

37-
# from core.NeuralLayers import *
54+
55+
# import torch
3856
# tensor_size = (1, 1, 32,32)
3957
# test = PointNet(tensor_size)
4058
# test(torch.rand(1,1,32,32))

0 commit comments

Comments
 (0)