Shortcuts

mmedit.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 mmedit.models.editors.indexnet.IndexNet(data_preprocessor, backbone, loss_alpha=None, loss_comp=None, init_cfg=None, train_cfg=None, test_cfg=None)[源代码]

Bases: mmedit.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 mmedit.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 mmedit.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 mmedit.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 mmedit.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 mmedit.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 nonlinearty 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

Read the Docs v: latest
Versions
master
latest
stable
zyh-doc-notfound-extend
Downloads
pdf
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.