206 lines
8.3 KiB
Python
206 lines
8.3 KiB
Python
# 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
|