"""Fully Convolutional Network with Strdie of 8"""
from __future__ import division
from mxnet.gluon import nn
import mxnet.ndarray as F
from mxnet.context import cpu
from mxnet.gluon.nn import HybridBlock
from .segbase import SegBaseModel
# pylint: disable=unused-argument,abstract-method,missing-docstring
__all__ = ['FCN', 'get_fcn', 'get_fcn_voc_resnet50', 'get_fcn_voc_resnet101']
[docs]class FCN(SegBaseModel):
r"""Fully Convolutional Networks for Semantic Segmentation
Parameters
----------
nclass : int
Number of categories for the training dataset.
backbone : string
Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
'resnet101' or 'resnet152').
norm_layer : object
Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`;
Reference:
Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks
for semantic segmentation." *CVPR*, 2015
Examples
--------
>>> model = FCN(nclass=21, backbone='resnet50')
>>> print(model)
"""
# pylint: disable=arguments-differ
def __init__(self, nclass, backbone='resnet50', norm_layer=nn.BatchNorm,
aux=True, **kwargs):
super(FCN, self).__init__(nclass, aux, backbone, norm_layer=norm_layer, **kwargs)
with self.name_scope():
self.head = _FCNHead(2048, nclass, norm_layer=norm_layer, **kwargs)
self.head.initialize()
self.head.collect_params().setattr('lr_mult', 10)
if self.aux:
self.auxlayer = _FCNHead(1024, nclass, norm_layer=norm_layer, **kwargs)
self.auxlayer.initialize()
self.auxlayer.collect_params().setattr('lr_mult', 10)
[docs] def forward(self, x):
_, _, H, W = x.shape
c3, c4 = self.base_forward(x)
outputs = []
x = self.head(c4)
x = F.contrib.BilinearResize2D(x, height=H, width=W)
outputs.append(x)
if self.aux:
auxout = self.auxlayer(c3)
auxout = F.contrib.BilinearResize2D(auxout, height=H, width=W)
outputs.append(auxout)
return tuple(outputs)
else:
return x
class _FCNHead(HybridBlock):
# pylint: disable=redefined-outer-name
def __init__(self, in_channels, channels, norm_layer, **kwargs):
super(_FCNHead, self).__init__()
with self.name_scope():
self.block = nn.HybridSequential()
inter_channels = in_channels // 4
with self.block.name_scope():
self.block.add(nn.Conv2D(in_channels=in_channels, channels=inter_channels,
kernel_size=3, padding=1))
self.block.add(norm_layer(in_channels=inter_channels))
self.block.add(nn.Activation('relu'))
self.block.add(nn.Dropout(0.1))
self.block.add(nn.Conv2D(in_channels=inter_channels, channels=channels,
kernel_size=1))
# pylint: disable=arguments-differ
def hybrid_forward(self, F, x):
return self.block(x)
[docs]def get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=False,
root='~/.mxnet/models', ctx=cpu(0), **kwargs):
r"""FCN model from the paper `"Fully Convolutional Network for semantic segmentation"
<https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_
Parameters
----------
dataset : str, default pascal_voc
The dataset that model pretrained on. (pascal_voc, ade20k)
pretrained : bool, default False
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '~/.mxnet/models'
Location for keeping the model parameters.
Examples
--------
>>> model = get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=False)
>>> print(model)
"""
acronyms = {
'pascal_voc': 'voc',
'ade20k': 'ade',
}
# infer number of classes
from ..data.segbase import get_segmentation_dataset
data = get_segmentation_dataset(dataset)
model = FCN(data.num_class, backbone=backbone, **kwargs)
if pretrained:
from .model_store import get_model_file
model.load_params(get_model_file('fcn_%s_%s'%(backbone, acronyms[dataset]),
root=root), ctx=ctx)
return model
[docs]def get_fcn_voc_resnet50(**kwargs):
r"""FCN model with base network ResNet-50 pre-trained on Pascal VOC dataset
from the paper `"Fully Convolutional Network for semantic segmentation"
<https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_
Parameters
----------
pretrained : bool, default False
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '~/.mxnet/models'
Location for keeping the model parameters.
Examples
--------
>>> model = get_fcn_voc_resnet50(pretrained=True)
>>> print(model)
"""
return get_fcn('pascal_voc', 'resnet50', **kwargs)
[docs]def get_fcn_voc_resnet101(**kwargs):
r"""FCN model with base network ResNet-101 pre-trained on Pascal VOC dataset
from the paper `"Fully Convolutional Network for semantic segmentation"
<https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_
Parameters
----------
pretrained : bool, default False
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '~/.mxnet/models'
Location for keeping the model parameters.
Examples
--------
>>> model = get_fcn_voc_resnet101(pretrained=True)
>>> print(model)
"""
return get_fcn('pascal_voc', 'resnet101', **kwargs)