mmagic.models.editors.indexnet.indexnet

Module Contents

Classes

IndexNet

IndexNet matting model.

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