def __init__(self,
in_channels,
inter_channels=None,
dimension=2,
sub_sample=True,
bn_layer=True,
is_training=True):
assert dimension in [1, 2, 3], "Unknown dimension..."
self.in_channels = in_channels
self.inter_channels = inter_channels
self.sub_sample = sub_sample
if self.inter_channels is None:
self.inter_channels = in_channels // 2
if self.inter_channels == 0:
self.inter_channels = 1
self.g = ConvBlock(self.inter_channels, dimension, 1)
self.W = ConvBlock(self.in_channels,
dimension,
1,
use_bn=bn_layer,
is_training=is_training)
if sub_sample:
self.g_pool = MaxPool(dimension)
self.phi = MaxPool(dimension)
def __init__(self, config):
super(Discriminator, self).__init__()
self.is_cuda = torch.cuda.is_available()
N = int(config.nfc)
self.head = ConvBlock(config.img_channel, N, config.kernel_size, 1, config.pad)
self.body = nn.Sequential()
for i in range(config.num_layers - 2):
N = int(config.nfc / pow(2, (i + 1)))
block = ConvBlock(max(2 * N, config.min_nfc), max(N, config.min_nfc), config.kernel_size, 1, config.pad)
self.body.add_module('block%d' % (i + 1), block)
# WGAN-GP discriminator has no activation at last layer
self.tail = nn.Conv2d(max(N, config.min_nfc), 1, kernel_size=config.kernel_size, stride=1, padding=config.pad)
def __init__(self, config, num_heads, real):
super(ACMDiscriminator, self).__init__()
self.ans = real
self.config = config
self.is_cuda = torch.cuda.is_available()
N = int(config.nfc)
self.head = ConvBlock(config.img_channel, N, config.kernel_size, 1, config.pad)
self.body = nn.Sequential()
for i in range(config.num_layers - 2):
N = int(config.nfc / pow(2, (i + 1)))
block = ConvBlock(max(2 * N, config.min_nfc), max(N, config.min_nfc), config.kernel_size, 1, config.pad)
self.body.add_module('block%d' % (i + 1), block)
self.acm = CustomACM(num_heads=num_heads, num_features=max(N, config.min_nfc), orthogonal_loss=self.config.use_acm_oth)
# WGAN-GP discriminator has no activation at last layer
self.tail = nn.Conv2d(max(N, config.min_nfc), 1, kernel_size=config.kernel_size, stride=1, padding=config.pad)
class DRNSeg(object):
def __init__(self,
model_name,
num_classes,
use_non_local=True,
use_bilinear=True,
is_training=True):
self.num_classes = num_classes
self.use_bilinear = use_bilinear
self.use_non_local = use_non_local
if model_name == 'drn_c_26':
self.base = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1],
num_classes=num_classes,
out_map=True,
out_middle=True,
is_training=is_training)
else:
raise Exception("Unknown architecture...")
if use_non_local:
self.nonLocal = NonLocalBlock2D(512)
self.seg = ConvBlock(num_classes, kernel_size=1)
# self.decoder = DecoderMudule(48, num_classes, use_bilinear=use_bilinear, is_training=is_training)
def forward(self, inputs):
inputs_shape = inputs.get_shape().as_list()
height, width = inputs_shape[1], inputs_shape[2]
_, end_points = self.base.forward(inputs)
out = end_points[-1]
encoder_features = end_points[1:3]
if self.use_non_local:
out = self.nonLocal(out)
out = self.seg.forward(out)
# out = self.decoder.forward(out, encoder_features)
if self.use_bilinear:
out = tf.image.resize_bilinear(out, (height, width))
else:
out = tf.layers.conv2d_transpose(
out,
self.num_classes,
2,
2,
padding='same',
use_bias=False,
kernel_initializer=tf.truncated_normal_initializer(
stddev=0.02))
pred = tf.argmax(out, axis=3, name='prediction')
pred = tf.expand_dims(pred, axis=3)
return out, pred
def __call__(self, inputs):
return self.forward(inputs)
class BasicBlock(object):
expansion = 1
def __init__(self,
out_channels,
stride=1,
dilation=(1, 1),
residual=True,
is_training=True):
self.out_channels = out_channels
self.stride = stride
self.residual = residual
self.is_training = is_training
self.conv_block1 = ConvBlock(out_channels,
stride=stride,
dilation=dilation[0],
use_bn=True,
use_act=True,
is_training=is_training)
self.conv_block2 = ConvBlock(out_channels,
dilation=dilation[1],
use_bn=True,
is_training=is_training)
def forward(self, inputs):
in_channels = tf.shape(inputs)[-1]
res = inputs
out = self.conv_block1.forward(inputs)
out = self.conv_block2.forward(out)
if self.stride != 1 or in_channels != self.out_channels:
downsample = ConvBlock(self.out_channels,
kernel_size=1,
stride=self.stride,
is_training=self.is_training)
res = downsample.forward(inputs)
if self.residual:
out += res
out = tf.nn.relu(out)
return out
def forward(self, inputs):
in_channels = tf.shape(inputs)[-1]
res = inputs
out = self.conv_block1.forward(inputs)
out = self.conv_block2.forward(out)
if self.stride != 1 or in_channels != self.out_channels:
downsample = ConvBlock(self.out_channels,
kernel_size=1,
stride=self.stride,
is_training=self.is_training)
res = downsample.forward(inputs)
if self.residual:
out += res
out = tf.nn.relu(out)
return out
def __init__(self,
out_channels,
stride=1,
dilation=(1, 1),
residual=True,
is_training=True):
self.out_channels = out_channels
self.stride = stride
self.residual = residual
self.is_training = is_training
self.conv_block1 = ConvBlock(out_channels,
stride=stride,
dilation=dilation[0],
use_bn=True,
use_act=True,
is_training=is_training)
self.conv_block2 = ConvBlock(out_channels,
dilation=dilation[1],
use_bn=True,
is_training=is_training)
def __init__(self, config):
super(Generator, self).__init__()
self.is_cuda = torch.cuda.is_available()
N = config.nfc
self.head = ConvBlock(config.img_channel, N, config.kernel_size, 1,
config.pad)
self.body = nn.Sequential()
for i in range(config.num_layers - 2):
N = int(config.nfc / pow(2, (i + 1)))
block = ConvBlock(max(2 * N,
config.min_nfc), max(N, config.min_nfc),
config.kernel_size, 1, config.pad)
self.body.add_module('block%d' % (i + 1), block)
self.tail = nn.Sequential(
nn.Conv2d(max(N, config.min_nfc),
config.img_channel,
kernel_size=config.kernel_size,
stride=1,
padding=config.pad), nn.Tanh())
def __init__(self,
model_name,
num_classes,
use_non_local=True,
use_bilinear=True,
is_training=True):
self.num_classes = num_classes
self.use_bilinear = use_bilinear
self.use_non_local = use_non_local
if model_name == 'drn_c_26':
self.base = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1],
num_classes=num_classes,
out_map=True,
out_middle=True,
is_training=is_training)
else:
raise Exception("Unknown architecture...")
if use_non_local:
self.nonLocal = NonLocalBlock2D(512)
self.seg = ConvBlock(num_classes, kernel_size=1)
def _build_layers(self):
self.conv1 = tf.make_template('conv1', ConvBlock(kernel_size=9,
filters=16,
data_format=self.data_format,
padding=self.padding))
self.conv2 = tf.make_template('conv2', ConvBlock(kernel_size=7,
filters=32,
data_format=self.data_format,
padding=self.padding))
self.conv3 = tf.make_template('conv3', ConvBlock(kernel_size=5,
filters=64,
data_format=self.data_format,
padding=self.padding))
self.conv4 = tf.make_template('conv4', ConvBlock(kernel_size=4,
filters=64,
data_format=self.data_format,
padding=self.padding))
# if use_edges is true, then the output from the Generator must be 3 channels, where:
# channel 0: probability of pixel is black i.e. background
# channel 1: probability pixel is white i.e. the body of a cell
# channel 2: 2, probability pixel is an edge of a cell
#
# if use_edges is false, then the output is a single channel where each value represents
# the probability of a pixel is a body of a cell
if self.use_edges:
self.conv5 = tf.make_template('conv5', ConvBlock(kernel_size=1,
filters=3,
data_format=self.data_format,
padding=self.padding,
is_final=True,
use_edges=True))
else:
self.conv5 = tf.make_template('conv5', ConvBlock(kernel_size=1,
filters=1,
data_format=self.data_format,
padding=self.padding,
is_final=True))
def __init__(self,
block,
layers,
num_classes=5,
channels=(16, 32, 64, 128, 256, 512, 512, 512),
out_map=False,
out_middle=False,
pool_size=32,
is_training=True):
self.out_map = out_map
self.pool_size = pool_size
self.out_middle = out_middle
self.conv1 = ConvBlock(channels[0],
kernel_size=7,
use_bias=False,
use_bn=True,
use_act=True,
is_training=is_training)
self.layer1 = MakeLayer(BasicBlock,
channels[0],
layers[0],
stride=1,
is_training=is_training)
self.layer2 = MakeLayer(BasicBlock,
channels[1],
layers[1],
stride=2,
is_training=is_training)
self.layer3 = MakeLayer(block,
channels[2],
layers[2],
stride=2,
is_training=is_training)
self.layer4 = MakeLayer(block,
channels[3],
layers[3],
stride=2,
is_training=is_training)
self.layer5 = MakeLayer(block,
channels[4],
layers[4],
dilation=2,
new_level=False,
is_training=is_training)
if layers[5] == 0:
self.layer6 = None
else:
self.layer6 = MakeLayer(block,
channels[5],
layers[5],
dilation=4,
new_level=False,
is_training=is_training)
if layers[6] == 0:
self.layer7 = None
else:
self.layer7 = MakeLayer(BasicBlock,
channels[6],
layers[6],
dilation=2,
new_level=False,
residual=False,
is_training=is_training)
if layers[7] == 0:
self.layer8 = None
else:
self.layer8 = MakeLayer(BasicBlock,
channels[7],
layers[7],
dilation=1,
new_level=False,
residual=False,
is_training=is_training)
if num_classes > 0:
self.fc = ConvBlock(num_classes,
kernel_size=1,
is_training=is_training)
class DRN(object):
def __init__(self,
block,
layers,
num_classes=5,
channels=(16, 32, 64, 128, 256, 512, 512, 512),
out_map=False,
out_middle=False,
pool_size=32,
is_training=True):
self.out_map = out_map
self.pool_size = pool_size
self.out_middle = out_middle
self.conv1 = ConvBlock(channels[0],
kernel_size=7,
use_bias=False,
use_bn=True,
use_act=True,
is_training=is_training)
self.layer1 = MakeLayer(BasicBlock,
channels[0],
layers[0],
stride=1,
is_training=is_training)
self.layer2 = MakeLayer(BasicBlock,
channels[1],
layers[1],
stride=2,
is_training=is_training)
self.layer3 = MakeLayer(block,
channels[2],
layers[2],
stride=2,
is_training=is_training)
self.layer4 = MakeLayer(block,
channels[3],
layers[3],
stride=2,
is_training=is_training)
self.layer5 = MakeLayer(block,
channels[4],
layers[4],
dilation=2,
new_level=False,
is_training=is_training)
if layers[5] == 0:
self.layer6 = None
else:
self.layer6 = MakeLayer(block,
channels[5],
layers[5],
dilation=4,
new_level=False,
is_training=is_training)
if layers[6] == 0:
self.layer7 = None
else:
self.layer7 = MakeLayer(BasicBlock,
channels[6],
layers[6],
dilation=2,
new_level=False,
residual=False,
is_training=is_training)
if layers[7] == 0:
self.layer8 = None
else:
self.layer8 = MakeLayer(BasicBlock,
channels[7],
layers[7],
dilation=1,
new_level=False,
residual=False,
is_training=is_training)
if num_classes > 0:
self.fc = ConvBlock(num_classes,
kernel_size=1,
is_training=is_training)
def forward(self, inputs):
end_points = []
out = self.conv1.forward(inputs)
out = self.layer1.forward(out)
end_points.append(out)
out = self.layer2.forward(out)
end_points.append(out)
out = self.layer3.forward(out)
end_points.append(out)
out = self.layer4.forward(out)
end_points.append(out)
out = self.layer5.forward(out)
end_points.append(out)
if self.layer6 is not None:
out = self.layer6.forward(out)
end_points.append(out)
if self.layer7 is not None:
out = self.layer7.forward(out)
end_points.append(out)
if self.layer8 is not None:
out = self.layer8.forward(out)
end_points.append(out)
if self.out_map:
out = self.fc.forward(out)
else:
out = tf.layers.average_pooling2d(out, self.pool_size,
self.pool_size)
out = self.fc.forward(out)
out = tf.reshape(out, (out.shape[0], -1))
if self.out_middle:
return out, end_points
else:
return out