head_extractor/src/mmseg/utils/mask_classification.py

# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple

import torch
from mmcv.ops import point_sample
from mmengine.structures import InstanceData
from torch import Tensor

from mmseg.registry import TASK_UTILS
from mmseg.utils import ConfigType, SampleList


def seg_data_to_instance_data(ignore_index: int,
                              batch_data_samples: SampleList):
    """Convert the paradigm of ground truth from semantic segmentation to
    instance segmentation.

    Args:
        ignore_index (int): The label index to be ignored.
        batch_data_samples (List[SegDataSample]): The Data
            Samples. It usually includes information such as
            `gt_sem_seg`.

    Returns:
        tuple[Tensor]: A tuple contains two lists.
            - batch_gt_instances (List[InstanceData]): Batch of
                gt_instance. It usually includes ``labels``, each is
                unique ground truth label id of images, with
                shape (num_gt, ) and ``masks``, each is ground truth
                masks of each instances of a image, shape (num_gt, h, w).
            - batch_img_metas (List[Dict]): List of image meta information.
    """
    batch_gt_instances = []

    for data_sample in batch_data_samples:
        gt_sem_seg = data_sample.gt_sem_seg.data
        classes = torch.unique(
            gt_sem_seg,
            sorted=False,
            return_inverse=False,
            return_counts=False)

        # remove ignored region
        gt_labels = classes[classes != ignore_index]

        masks = []
        for class_id in gt_labels:
            masks.append(gt_sem_seg == class_id)

        if len(masks) == 0:
            gt_masks = torch.zeros(
                (0, gt_sem_seg.shape[-2],
                 gt_sem_seg.shape[-1])).to(gt_sem_seg).long()
        else:
            gt_masks = torch.stack(masks).squeeze(1).long()

        instance_data = InstanceData(labels=gt_labels, masks=gt_masks)
        batch_gt_instances.append(instance_data)
    return batch_gt_instances


class MatchMasks:
    """Match the predictions to category labels.

    Args:
        num_points (int): the number of sampled points to compute cost.
        num_queries (int): the number of prediction masks.
        num_classes (int): the number of classes.
        assigner (BaseAssigner): the assigner to compute matching.
    """

    def __init__(self,
                 num_points: int,
                 num_queries: int,
                 num_classes: int,
                 assigner: ConfigType = None):
        assert assigner is not None, "\'assigner\' in decode_head.train_cfg" \
                                     'cannot be None'
        assert num_points > 0, 'num_points should be a positive integer.'
        self.num_points = num_points
        self.num_queries = num_queries
        self.num_classes = num_classes
        self.assigner = TASK_UTILS.build(assigner)

    def get_targets(self, cls_scores: List[Tensor], mask_preds: List[Tensor],
                    batch_gt_instances: List[InstanceData]) -> Tuple:
        """Compute best mask matches for all images for a decoder layer.

        Args:
            cls_scores (List[Tensor]): Mask score logits from a single
                decoder layer for all images. Each with shape (num_queries,
                cls_out_channels).
            mask_preds (List[Tensor]): Mask logits from a single decoder
                layer for all images. Each with shape (num_queries, h, w).
            batch_gt_instances (List[InstanceData]): each contains
                ``labels`` and ``masks``.

        Returns:
            tuple: a tuple containing the following targets.

                - labels (List[Tensor]): Labels of all images.\
                    Each with shape (num_queries, ).
                - mask_targets (List[Tensor]): Mask targets of\
                    all images. Each with shape (num_queries, h, w).
                - mask_weights (List[Tensor]): Mask weights of\
                    all images. Each with shape (num_queries, ).
                - avg_factor (int): Average factor that is used to
                    average the loss. `avg_factor` is usually equal
                    to the number of positive priors.
        """
        batch_size = cls_scores.shape[0]
        results = dict({
            'labels': [],
            'mask_targets': [],
            'mask_weights': [],
        })
        for i in range(batch_size):
            labels, mask_targets, mask_weights\
                = self._get_targets_single(cls_scores[i],
                                           mask_preds[i],
                                           batch_gt_instances[i])
            results['labels'].append(labels)
            results['mask_targets'].append(mask_targets)
            results['mask_weights'].append(mask_weights)

        # shape (batch_size, num_queries)
        labels = torch.stack(results['labels'], dim=0)
        # shape (batch_size, num_gts, h, w)
        mask_targets = torch.cat(results['mask_targets'], dim=0)
        # shape (batch_size, num_queries)
        mask_weights = torch.stack(results['mask_weights'], dim=0)

        avg_factor = sum(
            [len(gt_instances.labels) for gt_instances in batch_gt_instances])

        res = (labels, mask_targets, mask_weights, avg_factor)

        return res

    def _get_targets_single(self, cls_score: Tensor, mask_pred: Tensor,
                            gt_instances: InstanceData) \
            -> Tuple[Tensor, Tensor, Tensor]:
        """Compute a set of best mask matches for one image.

        Args:
            cls_score (Tensor): Mask score logits from a single decoder layer
                for one image. Shape (num_queries, cls_out_channels).
            mask_pred (Tensor): Mask logits for a single decoder layer for one
                image. Shape (num_queries, h, w).
            gt_instances (:obj:`InstanceData`): It contains ``labels`` and
                ``masks``.

        Returns:
            tuple[Tensor]: A tuple containing the following for one image.

                - labels (Tensor): Labels of each image. \
                    shape (num_queries, ).
                - mask_targets (Tensor): Mask targets of each image. \
                    shape (num_queries, h, w).
                - mask_weights (Tensor): Mask weights of each image. \
                    shape (num_queries, ).
        """
        gt_labels = gt_instances.labels
        gt_masks = gt_instances.masks
        # when "gt_labels" is empty, classify all queries to background
        if len(gt_labels) == 0:
            labels = gt_labels.new_full((self.num_queries, ),
                                        self.num_classes,
                                        dtype=torch.long)
            mask_targets = gt_labels
            mask_weights = gt_labels.new_zeros((self.num_queries, ))
            return labels, mask_targets, mask_weights
        # sample points
        num_queries = cls_score.shape[0]
        num_gts = gt_labels.shape[0]

        point_coords = torch.rand((1, self.num_points, 2),
                                  device=cls_score.device)
        # shape (num_queries, num_points)
        mask_points_pred = point_sample(
            mask_pred.unsqueeze(1), point_coords.repeat(num_queries, 1,
                                                        1)).squeeze(1)
        # shape (num_gts, num_points)
        gt_points_masks = point_sample(
            gt_masks.unsqueeze(1).float(), point_coords.repeat(num_gts, 1,
                                                               1)).squeeze(1)

        sampled_gt_instances = InstanceData(
            labels=gt_labels, masks=gt_points_masks)
        sampled_pred_instances = InstanceData(
            scores=cls_score, masks=mask_points_pred)
        # assign and sample
        matched_quiery_inds, matched_label_inds = self.assigner.assign(
            pred_instances=sampled_pred_instances,
            gt_instances=sampled_gt_instances)
        labels = gt_labels.new_full((self.num_queries, ),
                                    self.num_classes,
                                    dtype=torch.long)
        labels[matched_quiery_inds] = gt_labels[matched_label_inds]

        mask_weights = gt_labels.new_zeros((self.num_queries, ))
        mask_weights[matched_quiery_inds] = 1
        mask_targets = gt_masks[matched_label_inds]

        return labels, mask_targets, mask_weights