Source code for mmedit.models.editors.gca.gca

# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional

import torch

from mmedit.models.base_models import BaseMattor
from mmedit.models.utils import get_unknown_tensor
from mmedit.registry import MODELS

[docs]class GCA(BaseMattor): """Guided Contextual Attention image matting model. Args: data_preprocessor (dict, optional): The pre-process config of :class:`BaseDataPreprocessor`. backbone (dict): Config of backbone. loss_alpha (dict): Config of the alpha prediction loss. Default: None. init_cfg (dict, optional): Initialization config dict. Default: None. train_cfg (dict): Config of training. In ``train_cfg``, ``train_backbone`` should be specified. If the model has a refiner, ``train_refiner`` should be specified. test_cfg (dict): Config of testing. In ``test_cfg``, If the model has a refiner, ``train_refiner`` should be specified. """ def __init__(self, data_preprocessor, backbone, loss_alpha=None, init_cfg: Optional[dict] = None, train_cfg=None, test_cfg=None): super().__init__( backbone=backbone, data_preprocessor=data_preprocessor, init_cfg=init_cfg, train_cfg=train_cfg, test_cfg=test_cfg) self.loss_alpha =
[docs] def _forward(self, inputs): """Forward function. Args: inputs (torch.Tensor): Input tensor. Returns: Tensor: Output tensor. """ raw_alpha = self.backbone(inputs) pred_alpha = (raw_alpha.tanh() + 1.0) / 2.0 return pred_alpha
[docs] def _forward_test(self, inputs): """Forward function for testing GCA model. Args: inputs (torch.Tensor): batch input tensor. Returns: Tensor: Output tensor of model. """ return self._forward(inputs)
[docs] def _forward_train(self, inputs, data_samples): """Forward function for training GCA model. Args: inputs (torch.Tensor): batch input tensor collated by :attr:`data_preprocessor`. data_samples (List[BaseDataElement]): data samples collated by :attr:`data_preprocessor`. Returns: dict: Contains the loss items and batch information. """ trimap = inputs[:, 3:, :, :] gt_alpha = torch.stack(tuple( for ds in data_samples)) pred_alpha = self._forward(inputs) # FormatTrimap(to_onehot=False) will change unknown_value to 1 # FormatTrimap(to_onehot=True) will shift to 3 dim, # get_unknown_tensor can handle that directly without knowing # unknown_value. weight = get_unknown_tensor(trimap, unknown_value=1) losses = {'loss': self.loss_alpha(pred_alpha, gt_alpha, weight)} return losses
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.