mmagic.models.editors.indexnet
¶
Package Contents¶
Classes¶
IndexNet matting model. |
|
Indexed upsample module. |
|
Decoder for IndexNet. |
|
Depthwise index block. |
|
Holistic Index Block. |
|
Encoder for IndexNet. |
- class mmagic.models.editors.indexnet.IndexNet(data_preprocessor, backbone, loss_alpha=None, loss_comp=None, init_cfg=None, train_cfg=None, test_cfg=None)[源代码]¶
Bases:
mmagic.models.base_models.BaseMattor
IndexNet matting model.
This implementation follows: Indices Matter: Learning to Index for Deep Image Matting
- 参数
data_preprocessor (dict, optional) – The pre-process config of
BaseDataPreprocessor
.backbone (dict) – Config of backbone.
train_cfg (dict) – Config of training. In ‘train_cfg’, ‘train_backbone’ should be specified.
test_cfg (dict) – Config of testing.
init_cfg (dict, optional) – The weight initialized config for
BaseModule
.loss_alpha (dict) – Config of the alpha prediction loss. Default: None.
loss_comp (dict) – Config of the composition loss. Default: None.
- _forward(inputs)¶
Forward function.
- 参数
inputs (torch.Tensor) – Input tensor.
- 返回
Output tensor.
- 返回类型
Tensor
- _forward_test(inputs)¶
Forward function for testing IndexNet model.
- 参数
inputs (torch.Tensor) – batch input tensor.
- 返回
Output tensor of model.
- 返回类型
Tensor
- _forward_train(inputs, data_samples)¶
Forward function for training IndexNet model.
- 参数
inputs (torch.Tensor) – batch input tensor collated by
data_preprocessor
.data_samples (List[BaseDataElement]) – data samples collated by
data_preprocessor
.
- 返回
Contains the loss items and batch information.
- 返回类型
dict
- class mmagic.models.editors.indexnet.IndexedUpsample(in_channels, out_channels, kernel_size=5, norm_cfg=dict(type='BN'), conv_module=ConvModule, init_cfg: Optional[dict] = None)[源代码]¶
Bases:
mmengine.model.BaseModule
Indexed upsample module.
- 参数
in_channels (int) – Input channels.
out_channels (int) – Output channels.
kernel_size (int, optional) – Kernel size of the convolution layer. Defaults to 5.
norm_cfg (dict, optional) – Config dict for normalization layer. Defaults to dict(type=’BN’).
conv_module (ConvModule | DepthwiseSeparableConvModule, optional) – Conv module. Defaults to ConvModule.
init_cfg (dict, optional) – Initialization config dict. Default: None.
- init_weights()¶
Init weights for the module.
- forward(x, shortcut, dec_idx_feat=None)¶
Forward function.
- 参数
x (Tensor) – Input feature map with shape (N, C, H, W).
shortcut (Tensor) – The shortcut connection with shape (N, C, H’, W’).
dec_idx_feat (Tensor, optional) – The decode index feature map with shape (N, C, H’, W’). Defaults to None.
- 返回
Output tensor with shape (N, C, H’, W’).
- 返回类型
Tensor
- class mmagic.models.editors.indexnet.IndexNetDecoder(in_channels, kernel_size=5, norm_cfg=dict(type='BN'), separable_conv=False, init_cfg: Optional[dict] = None)[源代码]¶
Bases:
mmengine.model.BaseModule
Decoder for IndexNet.
Please refer to https://arxiv.org/abs/1908.00672.
- 参数
in_channels (int) – Input channels of the decoder.
kernel_size (int, optional) – Kernel size of the convolution layer. Defaults to 5.
norm_cfg (None | dict, optional) – Config dict for normalization layer. Defaults to dict(type=’BN’).
separable_conv (bool) – Whether to use separable conv. Default: False.
init_cfg (dict, optional) – Initialization config dict. Default: None.
- init_weights()¶
Init weights for the module.
- forward(inputs)¶
Forward function.
- 参数
inputs (dict) – Output dict of IndexNetEncoder.
- 返回
Predicted alpha matte of the current batch.
- 返回类型
Tensor
- class mmagic.models.editors.indexnet.DepthwiseIndexBlock(in_channels, norm_cfg=dict(type='BN'), use_context=False, use_nonlinear=False, mode='o2o')[源代码]¶
Bases:
mmengine.model.BaseModule
Depthwise index block.
From https://arxiv.org/abs/1908.00672.
- 参数
in_channels (int) – Input channels of the holistic index block.
norm_cfg (dict) – Config dict for normalization layer. Default: dict(type=’BN’).
use_context (bool, optional) – Whether use larger kernel size in index block. Refer to the paper for more information. Defaults to False.
use_nonlinear (bool) – Whether add a non-linear conv layer in the index blocks. Default: False.
mode (str) – Mode of index block. Should be ‘o2o’ or ‘m2o’. In ‘o2o’ mode, the group of the conv layers is 1; In ‘m2o’ mode, the group of the conv layer is in_channels.
- forward(x)¶
Forward function.
- 参数
x (Tensor) – Input feature map with shape (N, C, H, W).
- 返回
Encoder index feature and decoder index feature.
- 返回类型
tuple(Tensor)
- class mmagic.models.editors.indexnet.HolisticIndexBlock(in_channels, norm_cfg=dict(type='BN'), use_context=False, use_nonlinear=False)[源代码]¶
Bases:
mmengine.model.BaseModule
Holistic Index Block.
From https://arxiv.org/abs/1908.00672.
- 参数
in_channels (int) – Input channels of the holistic index block.
norm_cfg (dict) – Config dict for normalization layer. Default: dict(type=’BN’).
use_context (bool, optional) – Whether use larger kernel size in index block. Refer to the paper for more information. Defaults to False.
use_nonlinear (bool) – Whether add a non-linear conv layer in the index block. Default: False.
- forward(x)¶
Forward function.
- 参数
x (Tensor) – Input feature map with shape (N, C, H, W).
- 返回
Encoder index feature and decoder index feature.
- 返回类型
tuple(Tensor)
- class mmagic.models.editors.indexnet.IndexNetEncoder(in_channels, out_stride=32, width_mult=1, index_mode='m2o', aspp=True, norm_cfg=dict(type='BN'), freeze_bn=False, use_nonlinear=True, use_context=True, init_cfg: Optional[dict] = None)[源代码]¶
Bases:
mmengine.model.BaseModule
Encoder for IndexNet.
Please refer to https://arxiv.org/abs/1908.00672.
- 参数
in_channels (int, optional) – Input channels of the encoder.
out_stride (int, optional) – Output stride of the encoder. For example, if out_stride is 32, the input feature map or image will be downsample to the 1/32 of original size. Defaults to 32.
width_mult (int, optional) – Width multiplication factor of channel dimension in MobileNetV2. Defaults to 1.
index_mode (str, optional) – Index mode of the index network. It must be one of {‘holistic’, ‘o2o’, ‘m2o’}. If it is set to ‘holistic’, then Holistic index network will be used as the index network. If it is set to ‘o2o’ (or ‘m2o’), when O2O (or M2O) Depthwise index network will be used as the index network. Defaults to ‘m2o’.
aspp (bool, optional) – Whether use ASPP module to augment output feature. Defaults to True.
norm_cfg (None | dict, optional) – Config dict for normalization layer. Defaults to dict(type=’BN’).
freeze_bn (bool, optional) – Whether freeze batch norm layer. Defaults to False.
use_nonlinear (bool, optional) – Whether use nonlinearity in index network. Refer to the paper for more information. Defaults to True.
use_context (bool, optional) – Whether use larger kernel size in index network. Refer to the paper for more information. Defaults to True.
init_cfg (dict, optional) – Initialization config dict. Default: None.
- 引发
ValueError – out_stride must 16 or 32.
NameError – Supported index_mode are {‘holistic’, ‘o2o’, ‘m2o’}.
- _make_layer(layer_setting, norm_cfg)¶
- train(mode=True)¶
Set BatchNorm modules in the model to evaluation mode.
- init_weights()¶
Init weights for the model.
Initialization is based on self._init_cfg
- 参数
pretrained (str, optional) – Path for pretrained weights. If given None, pretrained weights will not be loaded. Defaults to None.
- forward(x)¶
Forward function.
- 参数
x (Tensor) – Input feature map with shape (N, C, H, W).
- 返回
Output tensor, shortcut feature and decoder index feature.
- 返回类型
dict