tfm.vision.heads.MaskHead

Creates a mask head.

tfm.vision.heads.MaskHead(
    num_classes: int,
    upsample_factor: int = 2,
    num_convs: int = 4,
    num_filters: int = 256,
    use_separable_conv: bool = False,
    activation: str = 'relu',
    use_sync_bn: bool = False,
    norm_momentum: float = 0.99,
    norm_epsilon: float = 0.001,
    kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
    bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
    class_agnostic: bool = False,
    **kwargs
)

num_classes An int of the number of classes.
upsample_factor An int that indicates the upsample factor to generate the final predicted masks. It should be >= 1.
num_convs An int number that represents the number of the intermediate convolution layers before the mask prediction layers.
num_filters An int number that represents the number of filters of the intermediate convolution layers.
use_separable_conv A bool that indicates whether the separable convolution layers is used.
activation A str that indicates which activation is used, e.g. 'relu', 'swish', etc.
use_sync_bn A bool that indicates whether to use synchronized batch normalization across different replicas.
norm_momentum A float of normalization momentum for the moving average.
norm_epsilon A float added to variance to avoid dividing by zero.
kernel_regularizer A tf.keras.regularizers.Regularizer object for Conv2D. Default is None.
bias_regularizer A tf.keras.regularizers.Regularizer object for Conv2D.
class_agnostic A bool. If set, we use a single channel mask head that is shared between all classes.
**kwargs Additional keyword arguments to be passed.

Methods

call

View source

call(
    inputs: List[tf.Tensor], training: bool = None
)

Forward pass of mask branch for the Mask-RCNN model.

Args
inputs A list of two tensors where inputs[0]: A tf.Tensor of shape [batch_size, num_instances, roi_height, roi_width, roi_channels], representing the ROI features. inputs[1]: A tf.Tensor of shape [batch_size, num_instances], representing the classes of the ROIs.
training A bool indicating whether it is in training mode.

Returns
mask_outputs A tf.Tensor of shape [batch_size, num_instances, roi_height * upsample_factor, roi_width * upsample_factor], representing the mask predictions.