mmagic.models.editors.indexnet

Package Contents

Classes

IndexNet	IndexNet matting model.
IndexedUpsample	Indexed upsample module.
IndexNetDecoder	Decoder for IndexNet.
DepthwiseIndexBlock	Depthwise index block.
HolisticIndexBlock	Holistic Index Block.
IndexNetEncoder	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

Parameters

• 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.

Parameters

inputs (torch.Tensor) – Input tensor.

Returns

Output tensor.

Return type

Tensor

_forward_test(inputs)

Forward function for testing IndexNet model.

Parameters

inputs (torch.Tensor) – batch input tensor.

Returns

Output tensor of model.

Return type

Tensor

_forward_train(inputs, data_samples)

Forward function for training IndexNet model.

Parameters

• inputs (torch.Tensor) – batch input tensor collated by data_preprocessor.

• data_samples (List[BaseDataElement]) – data samples collated by data_preprocessor.

Returns

Contains the loss items and batch information.

Return type

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.

Parameters

• 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.

Parameters

• 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.

Returns

Output tensor with shape (N, C, H’, W’).

Return type

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.

Parameters

• 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.

Parameters

inputs (dict) – Output dict of IndexNetEncoder.

Returns

Predicted alpha matte of the current batch.

Return type

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.

Parameters

• 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.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Encoder index feature and decoder index feature.

Return type

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.

Parameters

• 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.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Encoder index feature and decoder index feature.

Return type

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.

Parameters

• 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.

Raises

• 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

Parameters

pretrained (str, optional) – Path for pretrained weights. If given None, pretrained weights will not be loaded. Defaults to None.

forward(x)

Forward function.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Output tensor, shortcut feature and decoder index feature.

Return type

dict