Create xdecoder_model.py

xdecoder/architectures/xdecoder_model.py

@@ -0,0 +1,622 @@
+# --------------------------------------------------------
+# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Xueyan Zou (xueyan@cs.wisc.edu)
+# --------------------------------------------------------
+
+import random
+from typing import Tuple
+from unicodedata import name
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+
+from .registry import register_model
+from ..utils import configurable
+from ..backbone import build_backbone, Backbone
+from ..body import build_xdecoder_head
+from ..modules import sem_seg_postprocess, bbox_postprocess
+from ..language import build_language_encoder
+from ..language.loss import vl_similarity
+
+from timm.models.layers import trunc_normal_
+from nltk.stem.lancaster import LancasterStemmer
+from detectron2.structures import Boxes, ImageList, Instances, BitMasks, BoxMode
+from detectron2.utils.memory import retry_if_cuda_oom
+from detectron2.data import MetadataCatalog
+from utils.misc import prompt_engineering
+
+st = LancasterStemmer()
+
+
+class X_Decoder_Model(nn.Module):
+    @configurable
+    def __init__(
+        self,
+        *,
+        backbone: Backbone,
+        sem_seg_head: nn.Module,
+        criterion: nn.Module,
+        losses: dict,
+        num_queries: int,
+        object_mask_threshold: float,
+        overlap_threshold: float,
+        metadata,
+        task_switch: dict,
+        phrase_prob: float,
+        size_divisibility: int,
+        sem_seg_postprocess_before_inference: bool,
+        pixel_mean: Tuple[float],
+        pixel_std: Tuple[float],
+        # inference
+        semantic_on: bool,
+        panoptic_on: bool,
+        instance_on: bool,
+        test_topk_per_image: int,
+        train_dataset_name: str,
+        retrieval_emsemble: bool,
+        backbone_dim: int,
+        dim_proj: int,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.sem_seg_head = sem_seg_head
+        self.criterion = criterion
+        self.losses = losses
+        self.num_queries = num_queries
+        self.overlap_threshold = overlap_threshold
+        self.object_mask_threshold = object_mask_threshold
+        self.metadata = metadata
+        if size_divisibility < 0:
+            # use backbone size_divisibility if not set
+            size_divisibility = self.backbone.size_divisibility
+        self.size_divisibility = size_divisibility
+        self.sem_seg_postprocess_before_inference = sem_seg_postprocess_before_inference
+        self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
+        self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
+
+        # additional args
+        self.semantic_on = semantic_on
+        self.instance_on = instance_on
+        self.panoptic_on = panoptic_on
+
+        # caption argument
+        self.task_switch = task_switch
+        self.phrase_prob = phrase_prob
+
+        self.test_topk_per_image = test_topk_per_image
+        self.train_class_names = None
+
+        self.retrieval_emsemble = retrieval_emsemble
+        # backbone itc loss
+        if task_switch['retrieval'] and retrieval_emsemble:
+            self.backbone_proj = nn.Parameter(torch.empty(backbone_dim, dim_proj))
+            trunc_normal_(self.backbone_proj, std=.02)
+
+        if not self.semantic_on:
+            assert self.sem_seg_postprocess_before_inference
+
+    @classmethod
+    def from_config(cls, cfg):
+        enc_cfg = cfg['MODEL']['ENCODER']
+        dec_cfg = cfg['MODEL']['DECODER']
+
+        task_switch = {'bbox': dec_cfg.get('DETECTION', False),
+                       'mask': dec_cfg.get('MASK', True),
+                       'caption': dec_cfg['CAPTION'].get('ENABLED', False),
+                       'captioning': dec_cfg['CAPTIONING'].get('ENABLED', False),
+                       'retrieval': dec_cfg['RETRIEVAL'].get('ENABLED', False),
+                       'grounding': dec_cfg['GROUNDING'].get('ENABLED', False)}
+
+        # build model
+        extra = {'task_switch': task_switch}
+        backbone = build_backbone(cfg)
+        lang_encoder = build_language_encoder(cfg)        
+        sem_seg_head = build_xdecoder_head(cfg, backbone.output_shape(), lang_encoder, extra)
+
+        # Training Settings.
+        loss_weights = {}
+        matcher = None
+        losses = {}
+        weight_dict = {}        
+        grd_weight = {}
+        top_x_layers = {}
+        criterion = None
+        train_dataset_name = None
+        phrase_prob = None
+        # Loss parameters:
+        deep_supervision = None
+        no_object_weight = None
+
+        return {
+            "backbone": backbone,
+            "sem_seg_head": sem_seg_head,
+            "criterion": criterion,
+            "losses": losses,
+            "num_queries": dec_cfg['NUM_OBJECT_QUERIES'],
+            "object_mask_threshold": dec_cfg['TEST']['OBJECT_MASK_THRESHOLD'],
+            "overlap_threshold": dec_cfg['TEST']['OVERLAP_THRESHOLD'],
+            "metadata": None,
+            "size_divisibility": dec_cfg['SIZE_DIVISIBILITY'],
+            "sem_seg_postprocess_before_inference": (
+                dec_cfg['TEST']['SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE']
+                or dec_cfg['TEST']['PANOPTIC_ON']
+                or dec_cfg['TEST']['INSTANCE_ON']
+            ),
+            "pixel_mean": cfg['INPUT']['PIXEL_MEAN'],
+            "pixel_std": cfg['INPUT']['PIXEL_STD'],
+            "task_switch": task_switch,
+            "phrase_prob": phrase_prob,
+            # inference
+            "semantic_on": dec_cfg['TEST']['SEMANTIC_ON'],
+            "instance_on": dec_cfg['TEST']['INSTANCE_ON'],
+            "panoptic_on": dec_cfg['TEST']['PANOPTIC_ON'],
+            "test_topk_per_image": cfg['MODEL']['DECODER']['TEST']['DETECTIONS_PER_IMAGE'],
+            "train_dataset_name": train_dataset_name,
+            "retrieval_emsemble": dec_cfg['RETRIEVAL']['ENSEMBLE'],
+            "backbone_dim": cfg['MODEL']['BACKBONE_DIM'],
+            "dim_proj": cfg['MODEL']['DIM_PROJ'],
+        }
+
+    @property
+    def device(self):
+        return self.pixel_mean.device
+
+    def forward(self, batched_inputs, mode=None):
+        if self.training:
+            assert False, "Not support trianing mode."
+        else:
+            if mode == 'retrieval':
+                return self.evaluate_retrieval(batched_inputs)
+            elif mode == 'captioning':
+                return self.evaluate_captioning(batched_inputs)
+            elif mode == 'classification':
+                return self.evaluate_classification(batched_inputs)
+            elif mode in ['grounding_phrasecut', 'grounding_refcoco']:
+                return self.evaluate_grounding(batched_inputs, mode)
+            else:
+                return self.evaluate(batched_inputs)
+        
+    def evaluate(self, batched_inputs):
+        images = [x["image"].to(self.device) for x in batched_inputs]
+        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+        
+        images = ImageList.from_tensors(images, self.size_divisibility)
+        img_bs = images.tensor.shape[0]
+
+        targets = targets_grounding = queries_grounding = None
+        features = self.backbone(images.tensor)
+        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
+
+        mask_cls_results = outputs["pred_logits"]
+        mask_pred_results = outputs["pred_masks"]
+        box_pred_results = outputs["pred_boxes"] if self.task_switch['bbox'] else [None for i in range(len(mask_pred_results))]
+        caption_pred_results = outputs["pred_captions"] if self.task_switch['caption'] else [None for i in range(len(mask_pred_results))]
+
+        # upsample masks
+        mask_pred_results = F.interpolate(
+            mask_pred_results,
+            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        input_size = mask_pred_results.shape[-2:]
+        keep_sem_bgd = self.metadata.keep_sem_bgd if hasattr(self.metadata, 'keep_sem_bgd') else False
+        del outputs
+
+        processed_results = []
+        for mask_cls_result, mask_pred_result, box_pred_result, caption_pred_result, input_per_image, image_size in zip(
+            mask_cls_results, mask_pred_results, box_pred_results, caption_pred_results, batched_inputs, images.image_sizes
+        ):
+            height = input_per_image.get("height", image_size[0])
+            width = input_per_image.get("width", image_size[1])
+            processed_results.append({})
+
+            if self.sem_seg_postprocess_before_inference:
+                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)
+
+            # semantic segmentation inference
+            if self.semantic_on:
+                r = retry_if_cuda_oom(self.semantic_inference)(mask_cls_result, mask_pred_result, keep_sem_bgd)
+                if not self.sem_seg_postprocess_before_inference:
+                    r = retry_if_cuda_oom(sem_seg_postprocess)(r, image_size, height, width)
+                processed_results[-1]["sem_seg"] = r
+
+            # panoptic segmentation inference
+            if self.panoptic_on:
+                panoptic_r = retry_if_cuda_oom(self.panoptic_inference)(mask_cls_result, mask_pred_result)
+                processed_results[-1]["panoptic_seg"] = panoptic_r
+            
+            # instance segmentation inference
+            if self.instance_on:
+                if self.task_switch['bbox']:
+                    box_pred_result = bbox_postprocess(box_pred_result, input_size, image_size, height, width)
+                instance_r = retry_if_cuda_oom(self.instance_inference)(mask_cls_result, mask_pred_result, box_pred_result)
+                processed_results[-1]["instances"] = instance_r
+            if self.task_switch['caption']:
+                processed_results[-1]["captions"] = caption_pred_result
+                processed_results[-1]["masks"] = mask_pred_result
+
+        return processed_results
+
+
+    def evaluate_retrieval(self, batched_inputs):
+        images = [x["image"].to(self.device) for x in batched_inputs]
+        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+        images = ImageList.from_tensors(images, self.size_divisibility)
+        img_bs = images.tensor.shape[0]
+        
+        targets = targets_grounding = queries_grounding = None
+        features = self.backbone(images.tensor)
+        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
+        v_emb_it = outputs['pred_captions'][:,-1]
+
+        # compute backbone score
+        if self.task_switch['retrieval'] and self.retrieval_emsemble:
+            _v_emb_it = features['res5']
+            bs,nc,_,_ = _v_emb_it.shape
+            _v_emb_it = _v_emb_it.reshape(bs,nc,-1)
+            _v_emb_it = F.adaptive_avg_pool1d(_v_emb_it, 1).reshape(bs,nc) @ self.backbone_proj
+
+        processed_results = []
+        for idx, batch_data in enumerate(batched_inputs):
+            caption_ids = []
+            t_emb_its = []
+            processed_results.append({})
+            for caption in batch_data['captions']:
+                lang_results = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings(caption)
+                t_emb_it = lang_results['class_emb']
+                caption_ids.append(batch_data['image_id'])
+                t_emb_its.append(t_emb_it)
+
+            t_emb_it = torch.cat(t_emb_its, dim=0)
+
+            image_embeds = [v_emb_it[idx].unsqueeze(0)]
+            if self.task_switch['retrieval'] and self.retrieval_emsemble:
+                image_embeds += [_v_emb_it[idx].unsqueeze(0)]
+            caption_results = {
+                    'image_embeds': image_embeds,
+                    'text_embeds': t_emb_it,
+                    'caption_ids': caption_ids,
+                    'image_ids': batch_data['image_id'],
+                }
+            processed_results[-1]["caption"] = caption_results            
+        return processed_results
+
+    def evaluate_captioning(self, batched_inputs, extra={}):
+        images = [x["image"].to(self.device) for x in batched_inputs]
+        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+        images = ImageList.from_tensors(images, self.size_divisibility)
+        img_bs = images.tensor.shape[0]
+
+        if not hasattr(self, 'start_token'):
+            self.start_token = torch.tensor([[49406]*77], device=self.device)
+        
+        targets = targets_grounding = queries_grounding = None
+        features = self.backbone(images.tensor)
+
+        captioning_mask = None
+        if 'captioning_mask' in batched_inputs[-1]:
+            captioning_mask = torch.cat([x['captioning_mask'] for x in batched_inputs])
+
+        extra.update({'start_token': self.start_token, 'captioning_mask': captioning_mask})
+        outputs = self.sem_seg_head(features, target_queries=queries_grounding, task='captioning_infer', extra=extra)
+
+        processed_results = []
+        for idx, batch_data in enumerate(batched_inputs):
+            processed_results.append({})
+            processed_results[-1]["captioning_token"] = outputs['pred_captionings'][idx]
+            processed_results[-1]["captioning_text"] = outputs['pred_texts'][idx].split('.')[0]
+            processed_results[-1]["image_id"] = batched_inputs[idx]['image_id']
+            
+        return processed_results
+
+    def evaluate_classification(self, batched_inputs):
+        images = [x["image"].to(self.device) for x in batched_inputs]
+        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+        images = ImageList.from_tensors(images, self.size_divisibility)
+        img_bs = images.tensor.shape[0]
+        
+        targets = targets_grounding = queries_grounding = None
+        features = self.backbone(images.tensor)
+        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
+
+        processed_results = []
+        for idx, batch_data in enumerate(batched_inputs):
+            processed_results.append({})
+            processed_results[-1]["pred_class"] = outputs['pred_logits'][idx,-1]
+        return processed_results
+
+    def evaluate_grounding_baseline(self, batched_inputs, mode):
+        images = [x["image"].to(self.device) for x in batched_inputs]
+        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+        images = ImageList.from_tensors(images, self.size_divisibility)
+        img_bs = images.tensor.shape[0]
+        
+        targets = targets_grounding = queries_grounding = None
+        features = self.backbone(images.tensor)
+        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
+
+        mask_pred_results = outputs["pred_masks"]
+        caption_pred_results = outputs["pred_captions"] if self.task_switch['caption'] else [None for i in range(len(mask_pred_results))]
+
+        # upsample masks
+        mask_pred_results = F.interpolate(
+            mask_pred_results,
+            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        processed_results = []
+        for mask_pred_result, caption_pred_result, input_per_image, image_size in zip(
+            mask_pred_results, caption_pred_results, batched_inputs, images.image_sizes
+        ):
+            height = input_per_image.get("height", image_size[0])
+            width = input_per_image.get("width", image_size[1])
+            processed_results.append({})
+
+            mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
+                mask_pred_result, image_size, height, width
+            )[:-1]
+
+            texts_all = input_per_image['groundings']['texts']
+            grd_masks = []
+            for texts in texts_all:
+                if mode == 'grounding_refcoco':
+                    self.sem_seg_head.predictor.lang_encoder.get_text_embeddings(texts, name='grounding', prompt=False, is_eval=True)
+                elif mode == 'grounding_phrasecut':
+                    self.sem_seg_head.predictor.lang_encoder.get_text_embeddings(texts, name='grounding', prompt=True, is_eval=False)
+                t_emb = getattr(self.sem_seg_head.predictor.lang_encoder, "{}_text_embeddings".format('grounding')).t()
+                v_emb = caption_pred_result[:-1]
+                v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
+                vt_sim = v_emb @ t_emb
+                max_id = vt_sim.max(0)[1][0]
+                grd_masks += [mask_pred_result[max_id]]
+            processed_results[-1]['grounding_mask'] = torch.stack(grd_masks)
+
+        return processed_results
+
+    def evaluate_grounding(self, batched_inputs, mode):
+        images = [x["image"].to(self.device) for x in batched_inputs]
+        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+        images = ImageList.from_tensors(images, self.size_divisibility)
+
+        extra = {}
+        # mask_pred_results = []
+        # for idx, batch_per_image in enumerate(batched_inputs):
+        #     grd_texts = batch_per_image['groundings']['texts']
+        #     grd_masks = []
+        #     for anno_text in grd_texts:
+        #         gtext = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings([anno_text[0]], name='grounding', token=False, norm=False)
+        #         token_emb = gtext['token_emb']
+        #         tokens = gtext['tokens']
+            
+        #         grd_emb = token_emb[0][tokens['attention_mask'].bool()[0]]
+        #         extra['grounding_tokens'] = grd_emb[:,None]
+
+        #         assert len(images.tensor) == 1, "grounding evaluation only support single batch size now"
+        #         features = self.backbone(images.tensor)
+        #         outputs = self.sem_seg_head(features, extra=extra, task='grounding_eval')
+                
+        #         pred_gmasks = outputs['pred_masks'][idx,self.num_queries:2*self.num_queries-1]
+        #         v_emb = outputs['pred_captions'][idx,self.num_queries:2*self.num_queries-1]
+        #         t_emb = grd_emb[-1:]
+
+        #         t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
+        #         v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)            
+
+        #         temperature = self.sem_seg_head.predictor.lang_encoder.logit_scale
+        #         out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
+                
+        #         matched_id = out_prob.max(0)[1]
+        #         grd_masks += [pred_gmasks[matched_id,:,:]]
+        #     mask_pred_results += [torch.cat(grd_masks)]
+
+        # comment for multi object inference.
+        mask_pred_results = []
+        for idx, batch_per_image in enumerate(batched_inputs):
+            grd_texts = batch_per_image['groundings']['texts']
+            grd_texts = [x[0] for x in grd_texts]
+
+            gtext = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings(grd_texts, name='grounding', token=False, norm=False)
+            token_emb = gtext['token_emb']
+            tokens = gtext['tokens']
+            query_emb = token_emb[tokens['attention_mask'].bool()]
+            extra['grounding_tokens'] = query_emb[:,None]
+
+            features = self.backbone(images.tensor)
+            outputs = self.sem_seg_head(features, extra=extra, task='grounding_eval')
+
+            pred_gmasks = outputs['pred_masks'][idx,self.num_queries:2*self.num_queries-1]
+            v_emb = outputs['pred_captions'][idx,self.num_queries:2*self.num_queries-1]
+            t_emb = gtext['class_emb']
+
+            t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
+            v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)            
+
+            temperature = self.sem_seg_head.predictor.lang_encoder.logit_scale
+            out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
+            
+            matched_id = out_prob.max(0)[1]
+            mask_pred_results += [pred_gmasks[matched_id,:,:]]
+
+        for i in range(len(mask_pred_results)):
+            # upsample masks
+            mask_pred_results[i] = F.interpolate(
+                mask_pred_results[i][None,],
+                size=(images.tensor.shape[-2], images.tensor.shape[-1]),
+                mode="bilinear",
+                align_corners=False,
+            )[0]
+
+        processed_results = []
+        for mask_pred_result, input_per_image, image_size in zip(
+            mask_pred_results, batched_inputs, images.image_sizes
+        ):
+            height = input_per_image.get("height", image_size[0])
+            width = input_per_image.get("width", image_size[1])
+            processed_results.append({})
+
+            mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
+                mask_pred_result, image_size, height, width
+            )
+            processed_results[-1]['grounding_mask'] = mask_pred_result
+
+            # compute bbox
+            # bbox = BitMasks(mask_pred_result > 0).get_bounding_boxes()
+            # bbox = BoxMode.convert(bbox.tensor, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+            # processed_results[-1]['grounding_box'] = bbox
+
+        return processed_results
+
+    def prepare_vlp_targets(self, batched_inputs, device):
+        input_ids = []
+        attention_mask = []
+        for cnt, x in enumerate(batched_inputs):
+            captions = x['captions']
+            randid = random.randint(0, len(captions)-1)
+            input_ids += x['tokens']['input_ids'][randid:randid+1]
+            attention_mask += x['tokens']['attention_mask'][randid:randid+1]
+
+        input_ids = torch.stack(input_ids)
+        attention_mask = torch.stack(attention_mask)
+        tokens = {"input_ids": input_ids, "attention_mask": attention_mask}
+        lang_results = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings(tokens, token=True)
+
+        target_vlp = []
+        for cnt, x in enumerate(batched_inputs):
+            target_dict = {}
+            target_dict["caption_tokens"] = lang_results['token_emb'][cnt:cnt+1]
+            target_dict["caption_proj"] = lang_results['class_emb'][cnt:cnt+1]
+            target_dict["caption_tokenids"] = lang_results['tokens']['input_ids'][cnt:cnt+1]
+            target_dict["caption_mask"] = lang_results['tokens']['attention_mask'][cnt:cnt+1]            
+            target_vlp.append(target_dict)
+        return target_vlp
+
+    def semantic_inference(self, mask_cls, mask_pred, keep_sem_bgd=False):
+        if keep_sem_bgd:
+            mask_cls = F.softmax(mask_cls, dim=-1)
+        else:
+            mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1]
+        mask_pred = mask_pred.sigmoid()
+        semseg = torch.einsum("qc,qhw->chw", mask_cls, mask_pred)
+        return semseg
+
+    def panoptic_inference(self, mask_cls, mask_pred):
+        scores, labels = F.softmax(mask_cls, dim=-1).max(-1)
+        mask_pred = mask_pred.sigmoid()
+
+        keep = labels.ne(self.sem_seg_head.num_classes) & (scores > self.object_mask_threshold)
+        cur_scores = scores[keep]
+        cur_classes = labels[keep]
+        cur_masks = mask_pred[keep]
+        cur_mask_cls = mask_cls[keep]
+        cur_mask_cls = cur_mask_cls[:, :-1]
+        cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks
+
+        h, w = cur_masks.shape[-2:]
+        panoptic_seg = torch.zeros((h, w), dtype=torch.int32, device=cur_masks.device)
+        segments_info = []
+
+        current_segment_id = 0
+
+        if cur_masks.shape[0] == 0:
+            # We didn't detect any mask :(
+            return panoptic_seg, segments_info
+        else:
+            # take argmax
+            cur_mask_ids = cur_prob_masks.argmax(0)
+            stuff_memory_list = {}
+            thing_dataset_id_to_contiguous_id = self.metadata.thing_dataset_id_to_contiguous_id if hasattr(self.metadata, 'thing_dataset_id_to_contiguous_id') else {}
+            for k in range(cur_classes.shape[0]):
+                pred_class = cur_classes[k].item()
+                isthing = pred_class in thing_dataset_id_to_contiguous_id.values()
+                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.overlap_threshold:
+                        continue
+
+                    # merge stuff regions
+                    if not isthing:
+                        if int(pred_class) in stuff_memory_list.keys():
+                            panoptic_seg[mask] = stuff_memory_list[int(pred_class)]
+                            continue
+                        else:
+                            stuff_memory_list[int(pred_class)] = current_segment_id + 1
+
+                    current_segment_id += 1
+                    panoptic_seg[mask] = current_segment_id
+
+                    segments_info.append(
+                        {
+                            "id": current_segment_id,
+                            "isthing": bool(isthing),
+                            "category_id": int(pred_class),
+                        }
+                    )
+            return panoptic_seg, segments_info
+
+    def instance_inference(self, mask_cls, mask_pred, box_pred):
+        # mask_pred is already processed to have the same shape as original input
+        image_size = mask_pred.shape[-2:]
+
+        # [Q, K]
+        scores = F.softmax(mask_cls, dim=-1)[:, :-1]
+        labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
+        # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
+        scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.test_topk_per_image, sorted=False)
+
+        labels_per_image = labels[topk_indices]
+        topk_indices = (topk_indices // self.sem_seg_head.num_classes)
+        # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
+        mask_pred = mask_pred[topk_indices]
+        if box_pred is not None:
+            box_pred = box_pred[topk_indices]
+
+        # if this is panoptic segmentation, we only keep the "thing" classes
+        if self.panoptic_on:
+            thing_dataset_id_to_contiguous_id = self.metadata.thing_dataset_id_to_contiguous_id if hasattr(self.metadata, 'thing_dataset_id_to_contiguous_id') else {}
+            keep = torch.zeros_like(scores_per_image).bool()
+            for i, lab in enumerate(labels_per_image):
+                keep[i] = lab in thing_dataset_id_to_contiguous_id.values()
+
+            scores_per_image = scores_per_image[keep]
+            labels_per_image = labels_per_image[keep]
+            mask_pred = mask_pred[keep]
+
+            if box_pred is not None:
+                box_pred = box_pred[keep]
+
+        result = Instances(image_size)
+        # mask (before sigmoid)
+        result.pred_masks = (mask_pred > 0).float()
+        # result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
+        # Uncomment the following to get boxes from masks (this is slow)
+
+        if box_pred is not None:
+            result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()
+        else:
+            result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
+
+        # calculate average mask prob
+        mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
+        result.scores = scores_per_image * mask_scores_per_image
+        result.pred_classes = labels_per_image
+
+        return result
+
+
+@register_model
+def get_segmentation_model(cfg, **kwargs):
+    return X_Decoder_Model(cfg)