mmdetection/projects/XDecoder/xdecoder/unified_head.py

import copy
from typing import Sequence

import torch
from mmengine.structures import InstanceData, PixelData
from torch import nn
from torch.nn import functional as F

from mmdet.evaluation.functional import INSTANCE_OFFSET
from mmdet.registry import MODELS
from .utils import (is_lower_torch_version, retry_if_cuda_oom,
                    sem_seg_postprocess)


@MODELS.register_module()
class XDecoderUnifiedhead(nn.Module):

    def __init__(self,
                 in_channels: int,
                 pixel_decoder: nn.Module,
                 transformer_decoder: nn.Module,
                 task: str = 'semseg',
                 test_cfg=None):
        super().__init__()
        self.task = task
        self.test_cfg = test_cfg

        pixel_decoder_ = copy.deepcopy(pixel_decoder)
        pixel_decoder_.update(in_channels=in_channels)
        self.pixel_decoder = MODELS.build(pixel_decoder_)

        transformer_decoder_ = copy.deepcopy(transformer_decoder)
        transformer_decoder_.update(task=task)
        self.predictor = MODELS.build(transformer_decoder_)

        self.return_inter_mask = False
        if self.task == 'ref-caption':
            # ref-caption = ref-seg + caption,
            # so we need to return the intermediate mask
            self.return_inter_mask = True

        self._all_text_prompts = None
        self._extra = None
        # TODO: Very trick, for retrieval task
        self._force_not_use_cache = False

    def pre_process(self, batch_data_samples, device):
        extra = {}
        if self.task != 'caption':
            # have text
            all_text_prompts = []
            num_thing_class = 0
            for data_samples in batch_data_samples:
                if isinstance(data_samples.text, str):
                    text = data_samples.text.split('.')
                elif isinstance(data_samples.text, Sequence):
                    text = data_samples.text
                else:
                    raise TypeError(
                        'Type pf data_sample.text must be sequence or str')
                text = list(filter(lambda x: len(x) > 0, text))
                all_text_prompts.append(text)
                num_thing_class = len(text)
                # for panoptic
                if 'stuff_text' in data_samples:
                    if isinstance(data_samples.stuff_text, str):
                        text = data_samples.stuff_text.split('.')
                    elif isinstance(data_samples.stuff_text, Sequence):
                        text = data_samples.stuff_text
                    else:
                        raise TypeError('Type pf data_sample.stuff_text '
                                        'must be sequence or str')
                    text = list(filter(lambda x: len(x) > 0, text))
                    all_text_prompts[-1].extend(text)

            # TODO: support batch
            all_text_prompts = all_text_prompts[0]

            if all_text_prompts != self._all_text_prompts \
                    or self._force_not_use_cache:
                # avoid redundant computation
                self._all_text_prompts = all_text_prompts
                if self.task in ['semseg', 'instance', 'panoptic']:
                    self.predictor.lang_encoder.get_mean_embeds(
                        all_text_prompts + ['background'])
                elif self.task == 'ref-seg':
                    token_info = self.predictor.lang_encoder.get_text_embeds(
                        all_text_prompts, norm=False)
                    token_emb = token_info['token_emb']
                    tokens = token_info['tokens']
                    query_emb = token_emb[tokens['attention_mask'].bool()]
                    extra['grounding_tokens'] = query_emb[:, None]
                    extra['class_emb'] = token_info['class_emb']
                elif self.task == 'retrieval':
                    token_info = self.predictor.lang_encoder.get_text_embeds(
                        all_text_prompts, norm=True)
                    extra['class_emb'] = token_info['class_emb']
                self._extra = extra
                return extra, all_text_prompts, num_thing_class
            else:
                return self._extra, all_text_prompts, num_thing_class
        else:
            if not hasattr(self, 'start_token'):
                self.start_token = self.predictor.lang_encoder. \
                    get_sot_token(device=device)
            extra['start_token'] = self.start_token
            return extra, None, None

    def predict(self, features, batch_data_samples):
        # multi scale feature
        mask_features, multi_scale_features = self.pixel_decoder(features)

        # pre process
        extra, all_text_prompts, num_thing_class = self.pre_process(
            batch_data_samples, mask_features.device)

        # transformer decoder forward
        predictions = self.predictor(
            multi_scale_features, mask_features, extra=extra)

        # post process
        return self.post_process(predictions, batch_data_samples,
                                 all_text_prompts, num_thing_class)

    def post_process(self, predictions, batch_data_samples, all_text_prompts,
                     num_thing_class):
        batch_img_metas = [
            data_samples.metainfo for data_samples in batch_data_samples
        ]
        batch_input_shape = batch_data_samples[0].metainfo['batch_input_shape']

        if self.task == 'caption':
            for text, data_samples in zip(predictions['pred_caption'],
                                          batch_data_samples):
                data_samples.pred_caption = text

            if 'pred_instances' in batch_data_samples[0]:
                for img_metas, data_samples in zip(batch_img_metas,
                                                   batch_data_samples):
                    original_caption = data_samples.text.split('.')
                    text_prompts = list(
                        filter(lambda x: len(x) > 0, original_caption))

                    height = img_metas['ori_shape'][0]
                    width = img_metas['ori_shape'][1]
                    image_size = img_metas['grounding_img_shape'][:2]

                    mask_pred_result = data_samples.pred_instances.masks.float(
                    )
                    mask_cls_result = data_samples.pred_instances.scores.float(
                    )

                    mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
                        mask_pred_result, image_size, height, width)

                    pred_instances = retry_if_cuda_oom(
                        self._instance_inference)(mask_cls_result,
                                                  mask_pred_result,
                                                  text_prompts)
                    data_samples.pred_instances = pred_instances

        elif self.task in ['semseg', 'instance', 'panoptic']:
            mask_pred_results = predictions['pred_masks']
            mask_cls_results = predictions['pred_logits']
            if is_lower_torch_version():
                mask_pred_results = F.interpolate(
                    mask_pred_results,
                    size=(batch_input_shape[-2], batch_input_shape[-1]),
                    mode='bicubic',
                    align_corners=False)
            else:
                mask_pred_results = F.interpolate(
                    mask_pred_results,
                    size=(batch_input_shape[-2], batch_input_shape[-1]),
                    mode='bicubic',
                    align_corners=False,
                    antialias=True)

            # for batch
            for mask_cls_result, \
                    mask_pred_result, \
                    img_metas, \
                    data_samples in zip(
                                mask_cls_results,
                                mask_pred_results,
                                batch_img_metas,
                                batch_data_samples):
                height = img_metas['ori_shape'][0]
                width = img_metas['ori_shape'][1]
                image_size = img_metas['img_shape'][:2]
                mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
                    mask_pred_result, image_size, height, width)
                mask_cls_result = mask_cls_result.to(mask_pred_result)

                if self.task == 'semseg':
                    pred_sem_seg = retry_if_cuda_oom(self._semantic_inference)(
                        mask_cls_result, mask_pred_result, all_text_prompts)
                    data_samples.pred_sem_seg = pred_sem_seg
                elif self.task == 'instance':
                    pred_instances = retry_if_cuda_oom(
                        self._instance_inference)(mask_cls_result,
                                                  mask_pred_result,
                                                  all_text_prompts)
                    data_samples.pred_instances = pred_instances
                elif self.task == 'panoptic':
                    pred_panoptic_seg = retry_if_cuda_oom(
                        self._panoptic_inference)(mask_cls_result,
                                                  mask_pred_result,
                                                  all_text_prompts,
                                                  num_thing_class)
                    data_samples.pred_panoptic_seg = pred_panoptic_seg
        elif self.task == 'ref-seg':
            mask_pred_results = predictions['pred_masks']
            mask_cls_results = predictions['pred_logits']
            results_ = zip(mask_pred_results, mask_cls_results,
                           batch_img_metas, batch_data_samples)
            for mask_pred_result, mask_cls_result, \
                    img_metas, data_samples in results_:
                if is_lower_torch_version():
                    mask_pred_result = F.interpolate(
                        mask_pred_result[None],
                        size=(batch_input_shape[-2], batch_input_shape[-1]),
                        mode='bicubic',
                        align_corners=False)[0]
                else:
                    mask_pred_result = F.interpolate(
                        mask_pred_result[None],
                        size=(batch_input_shape[-2], batch_input_shape[-1]),
                        mode='bicubic',
                        align_corners=False,
                        antialias=True)[0]

                if self.return_inter_mask:
                    mask = mask_pred_result > 0
                    pred_instances = InstanceData()
                    pred_instances.masks = mask
                    pred_instances.scores = mask_cls_result
                    data_samples.pred_instances = pred_instances
                    continue

                height = img_metas['ori_shape'][0]
                width = img_metas['ori_shape'][1]
                image_size = img_metas['img_shape'][:2]
                mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
                    mask_pred_result, image_size, height, width)

                pred_instances = retry_if_cuda_oom(self._instance_inference)(
                    mask_cls_result, mask_pred_result, all_text_prompts)
                data_samples.pred_instances = pred_instances
        elif self.task == 'retrieval':
            batch_data_samples[0].pred_score = predictions['pred_logits']
        return batch_data_samples

    def _instance_inference(self, mask_cls, mask_pred, text_prompts):
        num_class = len(text_prompts)

        if self.task in ['ref-seg', 'caption']:
            scores = F.softmax(mask_cls, dim=-1)
            scores_per_image = scores.max(dim=-1)[0]
            labels_per_image = torch.arange(num_class)
        else:
            scores = F.softmax(mask_cls, dim=-1)[:, :-1]

            labels = torch.arange(
                num_class,
                device=scores.device).unsqueeze(0).repeat(scores.shape[0],
                                                          1).flatten(0, 1)
            scores_per_image, topk_indices = scores.flatten(0, 1).topk(
                self.test_cfg.get('max_per_img', 100), sorted=False)

            labels_per_image = labels[topk_indices]
            topk_indices = (topk_indices // num_class)
            mask_pred = mask_pred[topk_indices]

        result = InstanceData()
        mask_pred = mask_pred.sigmoid()
        result.masks = (mask_pred > self.test_cfg.mask_thr).float()

        # calculate average mask prob
        mask_scores_per_image = (mask_pred.flatten(1) *
                                 result.masks.flatten(1)).sum(1) / (
                                     result.masks.flatten(1).sum(1) + 1e-6)
        result.scores = scores_per_image * mask_scores_per_image
        result.labels = labels_per_image
        result.label_names = [
            text_prompts[label] for label in labels_per_image
        ]
        result.bboxes = result.scores.new_zeros(len(result.scores), 4)
        return result

    def _semantic_inference(self, mask_cls, mask_pred, text_prompts):
        mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1]
        mask_pred = mask_pred.sigmoid()
        sem_seg = torch.einsum('qc,qhw->chw', mask_cls, mask_pred)

        if sem_seg.shape[0] == 1:
            # 0 is foreground, ignore_index is background
            sem_seg = (sem_seg.squeeze(0) <= self.test_cfg.mask_thr).int()
            sem_seg[sem_seg == 1] = self.test_cfg.get('ignore_index', 255)
        else:
            # 0 is foreground, ignore_index is background
            if self.test_cfg.use_thr_for_mc:
                foreground_flag = sem_seg > self.test_cfg.mask_thr
                sem_seg = sem_seg.max(0)[1]
                sem_seg[foreground_flag.sum(0) == 0] = self.test_cfg.get(
                    'ignore_index', 255)
            else:
                sem_seg = sem_seg.max(0)[1]
        pred_sem_seg = PixelData(
            sem_seg=sem_seg[None],
            metainfo={
                'label_names': text_prompts,
                'ignore_index': self.test_cfg.get('ignore_index', 255)
            })
        return pred_sem_seg

    def _panoptic_inference(self, mask_cls, mask_pred, all_text_prompts,
                            num_thing_class):
        scores, labels = F.softmax(mask_cls, dim=-1).max(-1)
        mask_pred = mask_pred.sigmoid()

        keep = labels.ne(len(all_text_prompts)) & (
            scores > self.test_cfg.mask_thr)
        cur_scores = scores[keep]
        cur_classes = labels[keep]
        cur_masks = mask_pred[keep]
        cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks

        h, w = cur_masks.shape[-2:]
        panoptic_seg = torch.full((h, w),
                                  self.test_cfg.get('ignore_index', 255),
                                  dtype=torch.int32,
                                  device=cur_masks.device)
        instance_id = 1

        if cur_masks.shape[0] > 0:
            cur_mask_ids = cur_prob_masks.argmax(0)
            for k in range(cur_classes.shape[0]):
                pred_class = cur_classes[k].item()
                isthing = int(pred_class) < num_thing_class
                mask_area = (cur_mask_ids == k).sum().item()
                original_area = (cur_masks[k] >= 0.5).sum().item()
                mask = (cur_mask_ids == k) & (cur_masks[k] >= 0.5)

                if mask_area > 0 and original_area > 0 and mask.sum().item(
                ) > 0:
                    if mask_area / original_area < self.test_cfg.overlap_thr:
                        continue
                    # merge stuff regions
                    if not isthing:
                        panoptic_seg[mask] = int(pred_class)
                    else:
                        panoptic_seg[mask] = int(
                            pred_class) + instance_id * INSTANCE_OFFSET
                        instance_id += 1

        panoptic_seg = PixelData(
            sem_seg=panoptic_seg[None],
            metainfo={
                'label_names': all_text_prompts,
                'ignore_index': self.test_cfg.get('ignore_index', 255)
            })
        return panoptic_seg