Skip to content

Commit

Permalink
Merge pull request #39 from Tensor46/develop
Browse files Browse the repository at this point in the history 
+ added anatomy net for segmentation
  • Loading branch information
@Tensor46
Tensor46 committed Jan 3, 2019
2 parents 7d3d63a + 099432a commit eace23c
Show file tree
Hide file tree
Showing 6 changed files with 397 additions and 185 deletions.
6 changes: 2 additions & 4 deletions core/NeuralArchitectures/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
"ResidualNet", "InceptionV4", "DenseNet",
"LinearVAE", "ConvolutionalVAE", "PGGAN",
"ContextNet", "FeatureNet", "FeatureCapNet",
"PointNet", "UNet", "UNetPatch", "UNetMini",
"PointNet", "UNet", "ANet",
"NeuralDecisionForest", "Models"]

from .capsulenet import CapsuleNet
Expand All @@ -20,15 +20,13 @@
from .convolutionalvae import ConvolutionalVAE
from .pggan import PGGAN
from .contextnet import ContextNet
from .featurenet import FeatureNet, FeatureCapNet
from .pointnet import PointNet
from .unet import UNet, UNetMini
from .unet import UNet, ANet
from .trees import NeuralDecisionForest

del trees
del unet
del pointnet
del featurenet
del contextnet
del pggan
del capsulenet
Expand Down
147 changes: 111 additions & 36 deletions core/NeuralArchitectures/contextnet.py
View file
Original file line number Diff line number Diff line change
@@ -1,51 +1,126 @@

""" TensorMONK's :: NeuralArchitectures """

import torch
""" TensorMONK's :: NeuralArchitectures """
import torch.nn as nn
import numpy as np
from ..NeuralLayers import *
#==============================================================================#
from ..NeuralLayers import Convolution
from ..NeuralLayers import ContextNet_Bottleneck


class ContextNet(nn.Module):
"""
Implemented https://arxiv.org/pdf/1805.04554.pdf
r"""ContextNet: Exploring Context and Detail for Semantic Seg in Real-time
Implemented from https://arxiv.org/pdf/1805.04554.pdf
Naming Convention for layers in the module:
cnv::Convolution:: regular convolution block
bn::bottleneck:: bottlenect residual block
dw::Convilution:: depth-wise seperable convolution block
dn:: deep netwrk for Context
sn:: shallow network for Context
Args:
tensor_size: shape of tensor in BCHW
"""
def __init__(self, tensor_size=(1, 3, 1024, 2048), *args, **kwargs):
super(ContextNet, self).__init__()
Strides = [2, 1, 1]
bottleneck = ContextNet_Bottleneck
normalization = "batch"
self.DeepNET = nn.Sequential()
normalization, strides = "batch", [2, 1, 1]
bottleneck = ContextNet_Bottleneck

self.DeepNET = nn.Sequential()
self.ShallowNET = nn.Sequential()
self.DeepNET.add_module("AVGPL", nn.AvgPool2d((5,5), (4,4), 2)) # 1, 1, 256, 512
self.DeepNET.add_module("DN_CNV1", Convolution(tensor_size, 3, 32, 2, True, "relu", 0., normalization, False, 1)) # 1, 1, 128, 256
self.DeepNET.add_module("DN_BN10", bottleneck(self.DeepNET[-1].tensor_size, 3, 32, 1, expansion=1))
self.DeepNET.add_module("DN_BN20", bottleneck(self.DeepNET[-1].tensor_size, 3, 32, 1, expansion=6)) # 1, 1, 128, 256
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
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
for i in range(2): self.DeepNET.add_module("DN_BN5"+str(i), bottleneck(self.DeepNET[-1].tensor_size, 3, 96, 1, expansion=6))
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
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
self.DeepNET.add_module("UPSMPLE", nn.Upsample(scale_factor = 4, mode = 'bilinear')) # 1, 1, 128, 256
_tensor_size = (1, 128, self.DeepNET[-2].tensor_size[2]*4, self.DeepNET[-2].tensor_size[3]*4)
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))
self.DeepNET.add_module("DN_DW12", Convolution(self.DeepNET[-1].tensor_size, 1, 128, 1, True, "relu", 0., normalization, False, 1))
self.DeepNET.add_module("DN_CNV3", Convolution(self.DeepNET[-1].tensor_size, 1, 128, 1, True, "relu", 0., normalization, False, 1)) # 128, 256

# 1, 1, 256, 512
self.DeepNET.add_module("avgpl", nn.AvgPool2d((5, 5), (4, 4), 2))
# 1, 1, 128, 256
self.DeepNET.add_module("dn_cnv1", Convolution(tensor_size, 3, 32, 2,
True, "relu", 0., normalization, False, 1))
self.DeepNET.add_module("dn_bn10",
bottleneck(self.DeepNET[-1].tensor_size, 3, 32,
1, expansion=1))
self.DeepNET.add_module("dn_bn20",
bottleneck(self.DeepNET[-1].tensor_size, 3, 32,
1, expansion=6)) # 1, 1, 128, 256
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
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
for i in range(2):
self.DeepNET.add_module("dn_bn5"+str(i),
bottleneck(self.DeepNET[-1].tensor_size, 3,
96, 1, expansion=6))
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
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
self.DeepNET.add_module("upsample", nn.Upsample(scale_factor=4,
mode='bilinear'))
# 1, 1, 128, 256
_tensor_size = (1, 128, self.DeepNET[-2].tensor_size[2]*4,
self.DeepNET[-2].tensor_size[3]*4)
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))
self.DeepNET.add_module("dn_dw12",
Convolution(self.DeepNET[-1].tensor_size, 1,
128, 1, True, "relu", 0.,
normalization, False, 1))
self.DeepNET.add_module("dn_cnv3",
Convolution(self.DeepNET[-1].tensor_size, 1,
128, 1, True, "relu", 0.,
normalization, False, 1))
# 128, 256
activation, pre_nm, groups = "relu", False, 1
self.ShallowNET.add_module("SM_CNV1", Convolution(tensor_size, 3, 32, 2, True, "relu", 0.,True, False, 1)) # 512 x 1024
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
self.ShallowNET.add_module("SM_DW12", Convolution(self.ShallowNET[-1].tensor_size, 1, 64, 1,True, activation, 0., normalization, pre_nm, groups))
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
self.ShallowNET.add_module("SM_DW22", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups))
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]))
self.ShallowNET.add_module("SM_DW32", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups))
self.ShallowNET.add_module("SM_CNV2", Convolution(self.ShallowNET[-1].tensor_size, 1, 128, 1,True, activation, 0., normalization, pre_nm, groups)) # 128, 256
self.ShallowNET.add_module("sm_cnv1",
Convolution(tensor_size, 3, 32, 2, True,
"relu", 0., True, False, 1))
# 512 x 1024
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
self.ShallowNET.add_module("sm_dw12",
Convolution(self.ShallowNET[-1].tensor_size,
1, 64, 1, True, activation, 0.,
normalization, pre_nm, groups))
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]))
self.ShallowNET.add_module("sm_dw22",
Convolution(self.ShallowNET[-1].tensor_size,
1, 128, 1, True, activation, 0.,
normalization, pre_nm, groups))
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]))
self.ShallowNET.add_module("sm_dw32",
Convolution(self.ShallowNET[-1].tensor_size,
1, 128, 1, True, activation, 0.,
normalization, pre_nm, groups))
self.ShallowNET.add_module("sm_cnv2",
Convolution(self.ShallowNET[-1].tensor_size,
1, 128, 1, True, activation, 0.,
normalization, pre_nm, groups))
# 128, 256
self.tensor_size = self.ShallowNET[-1].tensor_size

self.FuseNET = Convolution(self.ShallowNET[-1].tensor_size, 1, self.ShallowNET[-1].tensor_size[1], 1, True)
self.FuseNET = Convolution(self.ShallowNET[-1].tensor_size, 1,
self.ShallowNET[-1].tensor_size[1], 1, True)

def forward(self, tensor):
return self.FuseNET(self.DeepNET(tensor)+self.ShallowNET(tensor))
Expand Down
64 changes: 41 additions & 23 deletions core/NeuralArchitectures/pointnet.py
View file
Original file line number Diff line number Diff line change
@@ -1,40 +1,58 @@
""" TensorMONK's :: NeuralArchitectures """
import torch.nn as nn
from core.NeuralLayers import Convolution
# =========================================================================== #

""" TensorMONK's :: NeuralArchitectures """

import torch
import torch.nn as nn
import numpy as np
from ..NeuralLayers import *
#==============================================================================#
class PointNet(nn.Module):
r"""
Implemented from paper: Learning Discriminative and Transformation
Covariant Local Feature Detectors
Args:
tensor_size: shape of tensor in BCHW
(None/any integer >0, channels, height, width)
out_channels: depth of output feature channels.
activation: None/relu/relu6/lklu/elu/prelu/tanh/sigm/maxo/rmxo/swish
normalization: None/batch/group/instance/layer/pixelwise
"""
Implemented http:https://openaccess.thecvf.com/content_cvpr_2017/papers/Zhang_Learning_Discriminative_and_CVPR_2017_paper.pdf
"""
def __init__(self, tensor_size=(1, 1, 32, 32), out_channels=2, *args, **kwargs):
def __init__(self, tensor_size=(1, 1, 32, 32), out_channels=2,
*args, **kwargs):
super(PointNet, self).__init__()
normalization = "batch"
self.PointNET = nn.Sequential()
self.PointNET.add_module("CONV1", Convolution(tensor_size, 5, 32, 1, False, "relu", 0., None))
activation = "relu"
self.PointNET = nn.Sequential()
self.PointNET.add_module("CONV1",
Convolution(tensor_size, 5, 32, 1, False,
activation, 0., None))
self.PointNET.add_module("POOL1", nn.MaxPool2d(2))
_tensor_size = self.PointNET[-2].tensor_size
_tensor_size = (_tensor_size[0], _tensor_size[1], _tensor_size[2]//2, _tensor_size[3]//2)
#print(_tensor_size)
self.PointNET.add_module("CONV2", Convolution(_tensor_size, 3, 128,1, False, "relu", 0.,normalization))
_tensor_size = (_tensor_size[0], _tensor_size[1],
_tensor_size[2]//2, _tensor_size[3]//2)
self.PointNET.add_module("CONV2",
Convolution(_tensor_size, 3, 128, 1, False,
activation, 0., normalization))
self.PointNET.add_module("POOL2", nn.MaxPool2d(2))
_tensor_size = self.PointNET[-2].tensor_size
_tensor_size = (_tensor_size[0], _tensor_size[1], _tensor_size[2]//2, _tensor_size[3]//2)
#print(_tensor_size)
self.PointNET.add_module("CONV3", Convolution(_tensor_size, 3, 128, 1, False, "relu", 0.,normalization))
#print(self.PointNET[-1].tensor_size)
self.PointNET.add_module("CONV4", Convolution(self.PointNET[-1].tensor_size, 3, 256, 1, False, "relu", 0.,normalization))
#print(self.PointNET[-1].tensor_size)
self.PointNET.add_module("CONV5", Convolution(self.PointNET[-1].tensor_size, 2, out_channels, 1, False, "relu", 0.,None))
print(self.PointNET[-1].tensor_size)
_tensor_size = (_tensor_size[0], _tensor_size[1],
_tensor_size[2]//2, _tensor_size[3]//2)
self.PointNET.add_module("CONV3",
Convolution(_tensor_size, 3, 128, 1, False,
activation, 0., normalization))
self.PointNET.add_module("CONV4",
Convolution(self.PointNET[-1].tensor_size, 3,
256, 1, False, activation, 0.,
normalization))
self.PointNET.add_module("CONV5",
Convolution(self.PointNET[-1].tensor_size, 2,
out_channels, 1, False,
activation, 0., None))
self.tensor_size = self.PointNET[-1].tensor_size

def forward(self, tensor):
return self.PointNET(tensor).squeeze(2).squeeze(2)

# from core.NeuralLayers import *

# import torch
# tensor_size = (1, 1, 32,32)
# test = PointNet(tensor_size)
# test(torch.rand(1,1,32,32))
Loading

0 comments on commit eace23c

Please sign in to comment.