mmagic.models.editors.indexnet.indexnet_decoder

Module Contents

Classes

IndexedUpsample	Indexed upsample module.
IndexNetDecoder	Decoder for IndexNet.

class mmagic.models.editors.indexnet.indexnet_decoder.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.indexnet_decoder.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