frozenseg/frozenseg.py

from typing import Tuple
import torch
from torch import nn
from torch.nn import functional as F

from detectron2.config import configurable
from detectron2.data import MetadataCatalog
from detectron2.modeling import META_ARCH_REGISTRY, build_backbone, build_sem_seg_head
from detectron2.modeling.backbone import Backbone
from detectron2.modeling.postprocessing import sem_seg_postprocess
from detectron2.structures import Boxes, ImageList, Instances
from detectron2.utils.memory import retry_if_cuda_oom
from .modeling.criterion import SetCriterion
from .modeling.matcher import HungarianMatcher
from segment_anything.build_sam import sam_model_registry
from .modeling.transformer_decoder.frozenseg_transformer_decoder import MaskPooling, get_classification_logits
from segment_anything import sam_model_registry
import pickle
VILD_PROMPT = [
    "a photo of a {}.",
    "This is a photo of a {}",
    "There is a {} in the scene",
    "There is the {} in the scene",
    "a photo of a {} in the scene",
    "a photo of a small {}.",
    "a photo of a medium {}.",
    "a photo of a large {}.",
    "This is a photo of a small {}.",
    "This is a photo of a medium {}.",
    "This is a photo of a large {}.",
    "There is a small {} in the scene.",
    "There is a medium {} in the scene.",
    "There is a large {} in the scene.",
]


@META_ARCH_REGISTRY.register()
class FrozenSeg(nn.Module):
    """
    Main class for mask classification semantic segmentation architectures.
    """

    @configurable
    def __init__(
        self,
        *,
        backbone: Backbone,
        sem_seg_head: nn.Module,
        criterion: nn.Module,
        num_queries: int,
        object_mask_threshold: float,
        overlap_threshold: float,
        train_metadata,
        test_metadata,
        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,
        geometric_ensemble_alpha: float,
        geometric_ensemble_beta: float,
        ensemble_on_valid_mask: bool,
        # FrozenSeg
        sam_name: str,
        mask_pred_alpha: float,
        use_sam_masks: bool,
        recall_on: bool,
        pkl_sam_name: str,
    ):
        """
        Args:
            backbone: a backbone module, must follow detectron2's backbone interface
            sem_seg_head: a module that predicts semantic segmentation from backbone features
            criterion: a module that defines the loss
            num_queries: int, number of queries
            object_mask_threshold: float, threshold to filter query based on classification score
                for panoptic segmentation inference
            overlap_threshold: overlap threshold used in general inference for panoptic segmentation
            metadata: dataset meta, get `thing` and `stuff` category names for panoptic
                segmentation inference
            size_divisibility: Some backbones require the input height and width to be divisible by a
                specific integer. We can use this to override such requirement.
            sem_seg_postprocess_before_inference: whether to resize the prediction back
                to original input size before semantic segmentation inference or after.
                For high-resolution dataset like Mapillary, resizing predictions before
                inference will cause OOM error.
            pixel_mean, pixel_std: list or tuple with #channels element, representing
                the per-channel mean and std to be used to normalize the input image
            semantic_on: bool, whether to output semantic segmentation prediction
            instance_on: bool, whether to output instance segmentation prediction
            panoptic_on: bool, whether to output panoptic segmentation prediction
            test_topk_per_image: int, instance segmentation parameter, keep topk instances per image
        """
        super().__init__()
        self.backbone = backbone
        self.sem_seg_head = sem_seg_head
        self.criterion = criterion
        self.num_queries = num_queries
        self.overlap_threshold = overlap_threshold
        self.object_mask_threshold = object_mask_threshold
        self.train_metadata = train_metadata
        self.test_metadata = test_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
        self.test_topk_per_image = test_topk_per_image

        if not self.semantic_on:
            assert self.sem_seg_postprocess_before_inference

        self.mask_pooling = MaskPooling()
        self.geometric_ensemble_alpha = geometric_ensemble_alpha
        self.geometric_ensemble_beta = geometric_ensemble_beta
        self.ensemble_on_valid_mask = ensemble_on_valid_mask

        self.train_text_classifier = None
        self.test_text_classifier = None
        self.void_embedding = nn.Embedding(1, backbone.dim_latent) # use this for void

        _, self.train_num_templates, self.train_class_names = self.prepare_class_names_from_metadata(train_metadata, train_metadata)
        self.category_overlapping_mask, self.test_num_templates, self.test_class_names = self.prepare_class_names_from_metadata(test_metadata, train_metadata)
        # sam args
        sam_ckpt_path = {
            'vit_t': './pretrained_checkpoint/mobile_sam.pt',
            'vit_b': './pretrained_checkpoint/sam_vit_b_01ec64.pth',
            'vit_l': './pretrained_checkpoint/sam_vit_l_0b3195.pth',
            'vit_h': './pretrained_checkpoint/sam_vit_h_4b8939.pth',
        }
        self.sam = sam_model_registry[sam_name](checkpoint=sam_ckpt_path[sam_name])
        # freeze SAM
        for name, param in self.sam.named_parameters():
            if 'mask_decoder' in name:
                param.requires_grad = False
            else:
                param.requires_grad = False

        if not self.training:
            sam_mask_for_each_dataset={
                                    "openvocab_ade20k_full_sem_seg_val": 'ade20k_val',
                                    'openvocab_coco_2017_val_panoptic_with_sem_seg':'',
                                    'openvocab_pascal_ctx59_sem_seg_val': 'pc_val',
                                    'openvocab_pascal_ctx459_sem_seg_val': 'pc_val',
                                    'openvocab_pascal21_sem_seg_val':'pas_val',
                                    "bdd10k_val_sem_seg":'bdd_val',
                                    }
            self.sam_mask_dataset_name = sam_mask_for_each_dataset[self.test_metadata.name]
        self.counter = 0
        self.use_sam_masks = use_sam_masks
        self.mask_pred_alpha = mask_pred_alpha
        self.recall_on = recall_on
        self.pkl_sam_name = pkl_sam_name


    def prepare_class_names_from_metadata(self, metadata, train_metadata):
        def split_labels(x):
            res = []
            for x_ in x:
                x_ = x_.replace(', ', ',')
                x_ = x_.split(',') # there can be multiple synonyms for single class
                res.append(x_)
            return res
        # get text classifier
        try:
            class_names = split_labels(metadata.stuff_classes) # it includes both thing and stuff
            train_class_names = split_labels(train_metadata.stuff_classes)
        except:
            # this could be for insseg, where only thing_classes are available
            class_names = split_labels(metadata.thing_classes)
            train_class_names = split_labels(train_metadata.thing_classes)
        train_class_names = {l for label in train_class_names for l in label}
        category_overlapping_list = []
        for test_class_names in class_names:
            is_overlapping = not set(train_class_names).isdisjoint(set(test_class_names))
            category_overlapping_list.append(is_overlapping)
        category_overlapping_mask = torch.tensor(
            category_overlapping_list, dtype=torch.long)
        
        def fill_all_templates_ensemble(x_=''):
            res = []
            for x in x_:
                for template in VILD_PROMPT:
                    res.append(template.format(x))
            return res, len(res) // len(VILD_PROMPT)
       
        num_templates = []
        templated_class_names = []
        for x in class_names:
            templated_classes, templated_classes_num = fill_all_templates_ensemble(x)
            templated_class_names += templated_classes
            num_templates.append(templated_classes_num) # how many templates for current classes
        class_names = templated_class_names
        return category_overlapping_mask, num_templates, class_names

    def set_metadata(self, metadata):
        self.test_metadata = metadata
        self.category_overlapping_mask, self.test_num_templates, self.test_class_names = self.prepare_class_names_from_metadata(metadata, self.train_metadata)
        self.test_text_classifier = None
        return

    def get_text_classifier(self):
        if self.training:
            if self.train_text_classifier is None:
                text_classifier = []
                # this is needed to avoid oom, which may happen when num of class is large
                bs = 128
                for idx in range(0, len(self.train_class_names), bs):
                    text_classifier.append(self.backbone.get_text_classifier(self.train_class_names[idx:idx+bs], self.device).detach())
                text_classifier = torch.cat(text_classifier, dim=0)

                # average across templates and normalization.
                text_classifier /= text_classifier.norm(dim=-1, keepdim=True)
                text_classifier = text_classifier.reshape(text_classifier.shape[0]//len(VILD_PROMPT), len(VILD_PROMPT), text_classifier.shape[-1]).mean(1)
                text_classifier /= text_classifier.norm(dim=-1, keepdim=True)
                self.train_text_classifier = text_classifier
            return self.train_text_classifier, self.train_num_templates
        else:
            if self.test_text_classifier is None:
                text_classifier = []
                # this is needed to avoid oom, which may happen when num of class is large
                bs = 128
                for idx in range(0, len(self.test_class_names), bs):
                    text_classifier.append(self.backbone.get_text_classifier(self.test_class_names[idx:idx+bs], self.device).detach())
                text_classifier = torch.cat(text_classifier, dim=0)

                # average across templates and normalization.
                text_classifier /= text_classifier.norm(dim=-1, keepdim=True)
                text_classifier = text_classifier.reshape(text_classifier.shape[0]//len(VILD_PROMPT), len(VILD_PROMPT), text_classifier.shape[-1]).mean(1)
                text_classifier /= text_classifier.norm(dim=-1, keepdim=True)
                self.test_text_classifier = text_classifier
            return self.test_text_classifier, self.test_num_templates

    @classmethod
    def from_config(cls, cfg):
        backbone = build_backbone(cfg)
        sem_seg_head = build_sem_seg_head(cfg, backbone.output_shape())

        # Loss parameters:
        deep_supervision = cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION
        no_object_weight = cfg.MODEL.MASK_FORMER.NO_OBJECT_WEIGHT

        # loss weights
        class_weight = cfg.MODEL.MASK_FORMER.CLASS_WEIGHT
        dice_weight = cfg.MODEL.MASK_FORMER.DICE_WEIGHT
        mask_weight = cfg.MODEL.MASK_FORMER.MASK_WEIGHT

        # building criterion
        matcher = HungarianMatcher(
            cost_class=class_weight,
            cost_mask=mask_weight,
            cost_dice=dice_weight,
            num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
        )

        weight_dict = {"loss_ce": class_weight, "loss_mask": mask_weight, "loss_dice": dice_weight}

        if deep_supervision:
            dec_layers = cfg.MODEL.MASK_FORMER.DEC_LAYERS
            aux_weight_dict = {}
            for i in range(dec_layers - 1):
                aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
            weight_dict.update(aux_weight_dict)

        losses = ["labels", "masks"]

        criterion = SetCriterion(
            sem_seg_head.num_classes,
            matcher=matcher,
            weight_dict=weight_dict,
            eos_coef=no_object_weight,
            losses=losses,
            num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
            oversample_ratio=cfg.MODEL.MASK_FORMER.OVERSAMPLE_RATIO,
            importance_sample_ratio=cfg.MODEL.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO,
        )
        #sem_seg_postprocess_before_inference: for panoptic and instance 
        return {
            "backbone": backbone,
            "sem_seg_head": sem_seg_head,
            "criterion": criterion,
            "num_queries": cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES,
            "object_mask_threshold": cfg.MODEL.MASK_FORMER.TEST.OBJECT_MASK_THRESHOLD,
            "overlap_threshold": cfg.MODEL.MASK_FORMER.TEST.OVERLAP_THRESHOLD,
            "train_metadata": MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
            "test_metadata": MetadataCatalog.get(cfg.DATASETS.TEST[0]),
            "size_divisibility": cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY,
            "sem_seg_postprocess_before_inference": (
                cfg.MODEL.MASK_FORMER.TEST.SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE
                or cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON
                or cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON
                or cfg.MODEL.MASK_FORMER.TEST.RECALL_ON
            ),
            "pixel_mean": cfg.MODEL.PIXEL_MEAN,
            "pixel_std": cfg.MODEL.PIXEL_STD,
            # inference
            "semantic_on": cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON,
            "instance_on": cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON,
            "panoptic_on": cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON,
            "recall_on": cfg.MODEL.MASK_FORMER.TEST.RECALL_ON,
            "test_topk_per_image": cfg.TEST.DETECTIONS_PER_IMAGE,
            "geometric_ensemble_alpha": cfg.MODEL.FROZEN_SEG.GEOMETRIC_ENSEMBLE_ALPHA,
            "geometric_ensemble_beta": cfg.MODEL.FROZEN_SEG.GEOMETRIC_ENSEMBLE_BETA,
            "ensemble_on_valid_mask": cfg.MODEL.FROZEN_SEG.ENSEMBLE_ON_VALID_MASK,
            # FrozenSeg
            "sam_name": cfg.MODEL.SAM_NAME,
            "mask_pred_alpha": cfg.TEST.SAM_MASK_PRED_ALPHA,
            'use_sam_masks': cfg.TEST.USE_SAM_MASKS,
            "pkl_sam_name":  cfg.TEST.PKL_SAM_MODEL_NAME
        }

    @property
    def device(self):
        return self.pixel_mean.device

    def preprocess_wo_norm(self, x, resize=(512, 512)):
        x = x.float()
        x = F.interpolate(
              x.unsqueeze(0),
              size=resize,
              mode="bilinear",
              align_corners=False,
            )
        return x[0]

    def forward(self, batched_inputs):
        """
        Args:
            batched_inputs: a list, batched outputs of :class:`DatasetMapper`.
                Each item in the list contains the inputs for one image.
                For now, each item in the list is a dict that contains:
                   * "image": Tensor, image in (C, H, W) format.
                   * "instances": per-region ground truth
                   * Other information that's included in the original dicts, such as:
                     "height", "width" (int): the output resolution of the model (may be different
                     from input resolution), used in inference.
        Returns:
            list[dict]:
                each dict has the results for one image. The dict contains the following keys:
                * "sem_seg" if semantic_on
                * "panoptic_seg" if panoptic_on
                * "instances" if instance_on
                * "recall_seg" if recall_on
        """
        images = [x["image"].to(self.device) for x in batched_inputs] #raw images 3 1024 1024
        if self.sam is None:
            sam_embedding = None
        else:
            images_sam = [(x-self.sam.pixel_mean)/self.sam.pixel_std for x in images]
        if not self.training:
            images_sam = ImageList.from_tensors(images_sam, self.size_divisibility) 
            images_sam = images_sam.tensor.to(self.device)
            images_sam = torch.stack([self.preprocess_wo_norm(x, resize=(1024,1024)) for x in images_sam], dim=0)
        else:
            images_sam = torch.stack(images_sam, dim=0).to(self.device)
        last_embedding, interm_embeddings = self.sam.image_encoder(images_sam)
        sam_embedding = (last_embedding, interm_embeddings)

        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
        images = ImageList.from_tensors(images, self.size_divisibility) 
        features = self.backbone(images.tensor)

        text_classifier, num_templates = self.get_text_classifier()
        text_classifier = torch.cat([text_classifier, F.normalize(self.void_embedding.weight, dim=-1)], dim=0)
        features['text_classifier'] = text_classifier
        features['num_templates'] = num_templates
        features['sam_embedding'] = sam_embedding
        features['sam'] = self.sam
        outputs = self.sem_seg_head(features) 

        if self.training:
            if "instances" in batched_inputs[0]:
                gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
                targets = self.prepare_targets(gt_instances, images)
            else:
                targets = None
            # bipartite matching-based loss
            losses = self.criterion(outputs, targets)

            for k in list(losses.keys()):
                if k in self.criterion.weight_dict:
                    losses[k] *= self.criterion.weight_dict[k]
                else:
                    # remove this loss if not specified in `weight_dict`
                    losses.pop(k)
            return losses
        else:
            mask_cls_results = outputs["pred_logits"]
            mask_pred_results = outputs["pred_masks"]
            clip_feature = features["clip_vis_dense"]
            mask_for_pooling = F.interpolate(mask_pred_results, size=clip_feature.shape[-2:],
                                                mode='bilinear', align_corners=False)

            if "convnext" in self.backbone.model_name.lower():
                pooled_clip_feature = self.mask_pooling(clip_feature, mask_for_pooling) #mask>0
                pooled_clip_feature = self.backbone.visual_prediction_forward(pooled_clip_feature)
            elif "rn" in self.backbone.model_name.lower():
                pooled_clip_feature = self.backbone.visual_prediction_forward(clip_feature, mask_for_pooling)
            else:
                raise NotImplementedError
            
            out_vocab_cls_results = get_classification_logits(pooled_clip_feature, text_classifier, self.backbone.clip_model.logit_scale, num_templates)
            in_vocab_cls_results = mask_cls_results[..., :-1] # remove void
            out_vocab_cls_results = out_vocab_cls_results[..., :-1] # remove void
            out_vocab_cls_probs = out_vocab_cls_results.softmax(-1)
            in_vocab_cls_results = in_vocab_cls_results.softmax(-1)
            category_overlapping_mask = self.category_overlapping_mask.to(self.device)
           
            if self.ensemble_on_valid_mask:
                # Only include out_vocab cls results on masks with valid pixels
                # We empirically find that this is important to obtain reasonable AP/mIOU score with ResNet CLIP models
                valid_masking = (mask_for_pooling > 0).to(mask_for_pooling).sum(-1).sum(-1) > 0
                valid_masking = valid_masking.to(in_vocab_cls_results.dtype).unsqueeze(-1)
                alpha = torch.ones_like(in_vocab_cls_results) * self.geometric_ensemble_alpha
                beta = torch.ones_like(in_vocab_cls_results) * self.geometric_ensemble_beta
                alpha = alpha * valid_masking
                beta = beta * valid_masking
            else:
                alpha = self.geometric_ensemble_alpha
                beta = self.geometric_ensemble_beta

            cls_logits_seen = (
                (in_vocab_cls_results ** (1 - alpha) * out_vocab_cls_probs**alpha).log()
                * category_overlapping_mask
            )
            cls_logits_unseen = (
                (in_vocab_cls_results ** (1 - beta) * out_vocab_cls_probs**beta).log()
                * (1 - category_overlapping_mask)
            )
            cls_results = cls_logits_seen + cls_logits_unseen

            # This is used to filtering void predictions. 
            is_void_prob = F.softmax(mask_cls_results, dim=-1)[..., -1:]
            mask_cls_probs = torch.cat([
                cls_results.softmax(-1) * (1.0 - is_void_prob),
                is_void_prob], dim=-1)
            mask_cls_results = torch.log(mask_cls_probs + 1e-8)
            mask_pred_results = F.interpolate(
                mask_pred_results,
                size=(images.tensor.shape[-2], images.tensor.shape[-1]),
                mode="bilinear",
                align_corners=False,
            )
          
            del outputs
            processed_results = []
            for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
                mask_cls_results, 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({})
                if self.sem_seg_postprocess_before_inference: # panoptic on
                    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.use_sam_masks:
                    assert not self.training
                    img_id = input_per_image.get("image_id", None)
                    if img_id is None:
                        filename = input_per_image.get('file_name', None)
                        if filename is None:
                            assert NameError, 'No image_id or file_name in input_per_image'
                        elif filename is not None:
                            img_id = filename.split('/')[-1].split('.')[0]
                    with open(f'output/SAM_masks_pred/{self.pkl_sam_name}_{self.sam_mask_dataset_name}/{img_id}.pkl', 'rb') as f:
                        everything_mask = pickle.load(f)
                    
                    sam_mask_pred = [torch.from_numpy(mask['preds']).to(mask_cls_result.device) for mask in everything_mask]
                    if len(sam_mask_pred) == 0:
                        sam_mask_pred = None
                        sam_cls_results = None
                        sam_iou_scores = None
                    else:
                        sam_mask_pred = torch.stack(sam_mask_pred, dim=0) # M, H, W
                        sam_iou_scores = [torch.tensor(mask['predicted_iou']).sigmoid().to(mask_cls_result.device) for mask in everything_mask]
                        sam_iou_scores = torch.stack(sam_iou_scores, dim=0) # M, 1
                        sam_mask_for_pooling_clip = F.interpolate(sam_mask_pred.unsqueeze(0), size=clip_feature.shape[-2:], mode="nearest")
                        sam_mask_for_pooling_clip = ImageList.from_tensors([sam_mask_for_pooling_clip[0]], self.size_divisibility)
                        if 'convnext' in self.backbone.model_name.lower():
                            sam_pooled_clip_feature = self.mask_pooling(clip_feature, sam_mask_for_pooling_clip.tensor.to(torch.float32))
                            sam_pooled_clip_feature = self.backbone.visual_prediction_forward(sam_pooled_clip_feature)
                            sam_cls_results = get_classification_logits(sam_pooled_clip_feature, text_classifier, self.backbone.clip_model.logit_scale, num_templates)
                        elif "rn" in self.backbone.model_name.lower():
                            sam_pooled_clip_feature = self.backbone.visual_prediction_forward(clip_feature, sam_mask_for_pooling_clip.tensor.to(torch.float32))
                            sam_cls_results = get_classification_logits(sam_pooled_clip_feature, text_classifier, self.backbone.clip_model.logit_scale, num_templates) # have nan
                        else:
                            print("not support")
                            raise NotImplementedError
                        sam_mask_pred = sam_mask_pred.to(mask_pred_result)
                        sam_cls_results = sam_cls_results.to(mask_cls_result)
                    if not self.sem_seg_postprocess_before_inference: #### For semantic segmentation and recall inference
                        mask_pred_result = mask_pred_result[:, :image_size[0], :image_size[1]]
                    del everything_mask
                #######################################################

                if self.recall_on and not self.use_sam_masks:
                    res = retry_if_cuda_oom(self.recall_inference)(mask_pred_result)
                    processed_results[-1]["recall_seg"] = res
                elif self.recall_on and self.use_sam_masks:
                    if sam_mask_pred.shape[-2:] != mask_pred_result.shape[-2:]: 
                        sam_mask_pred = F.interpolate(sam_mask_pred.unsqueeze(0), size=mask_pred_result.shape[-2:], mode="bilinear", align_corners=False)[0]
                    res = retry_if_cuda_oom(self.recall_inference_with_everything)(mask_pred_result, sam_mask_pred)
                    res = retry_if_cuda_oom(sem_seg_postprocess)(res, image_size, height, width)
                    processed_results[-1]["recall_seg"] = res

                if self.semantic_on:
                    if self.use_sam_masks:
                        if sam_mask_pred.shape[-2:] != mask_pred_result.shape[-2:]: 
                            sam_mask_pred = F.interpolate(sam_mask_pred.unsqueeze(0), size=mask_pred_result.shape[-2:], mode="bilinear", align_corners=False)[0]
                        res = self.geo_with_sam_inference(mask_cls_result, mask_pred_result, sam_mask_pred, sam_cls_results, category_overlapping_mask)
                    else:
                        res = retry_if_cuda_oom(self.semantic_inference)(mask_cls_result, mask_pred_result)
                    if not self.sem_seg_postprocess_before_inference : # for sem seg
                        res = retry_if_cuda_oom(sem_seg_postprocess)(res, image_size, height, width)
                    processed_results[-1]["sem_seg"] = res
                
                # 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:
                    instance_r = retry_if_cuda_oom(self.instance_inference)(mask_cls_result, mask_pred_result)
                    processed_results[-1]["instances"] = instance_r
            return processed_results

    def prepare_targets(self, targets, images):
        h_pad, w_pad = images.tensor.shape[-2:]
        new_targets = []
        for targets_per_image in targets:
            # pad gt
            gt_masks = targets_per_image.gt_masks
            padded_masks = torch.zeros((gt_masks.shape[0], h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device)
            padded_masks[:, : gt_masks.shape[1], : gt_masks.shape[2]] = gt_masks
            new_targets.append(
                {
                    "labels": targets_per_image.gt_classes,
                    "masks": padded_masks,
                }
            )
        return new_targets
    

    def resize_feat(self, x, resize_shape):
        x = F.interpolate(
              x,
              size=(resize_shape[0], resize_shape[1]),
              mode="bilinear",
              align_corners=False,
            )
        return x
    
    def recall_inference(self, mask_pred):
        """
        Return: (q, h, w)
        """
        return mask_pred

    def recall_inference_with_everything(self, mask_pred, sam_mask_pred):
        """
        Return: (q, h, w)
        """
        if sam_mask_pred is None:
            return self.recall_inference(mask_pred)
        return torch.cat([mask_pred, sam_mask_pred], dim=0)

    def semantic_inference(self, mask_cls, mask_pred):
        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()
        num_classes = len(self.test_metadata.stuff_classes)
        keep = labels.ne(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 = {}
            for k in range(cur_classes.shape[0]): # through all mask queries
                pred_class = cur_classes[k].item()
                isthing = pred_class in self.test_metadata.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: # 0.8 for coco, 0 for else.
                        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 geo_with_sam_inference(self, mask_cls, mask_pred, sam_mask_pred, sam_mask_cls, category_overlapping_mask=None):
        if sam_mask_cls is None:
            return self.semantic_inference(mask_cls, mask_pred)

        mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1] 
        sam_mask_cls = F.softmax(sam_mask_cls, dim=-1)[..., :-1].squeeze(0) # M, C
        mask_pred = mask_pred.sigmoid()
        sam_mask_pred = sam_mask_pred.sigmoid()

        alpha = self.mask_pred_alpha
        beta = 0.
        in_mask_cls = mask_cls * category_overlapping_mask.view(1,-1) # (q,c)
        out_mask_cls = mask_cls * (1 - category_overlapping_mask).view(1,-1)  # (q,c) # 0
        out_mask_cls = out_mask_cls * (1-alpha)
        in_mask_cls = in_mask_cls * (1-beta)

        ## MaskEnsemble
        in_sam_mask_cls = sam_mask_cls * category_overlapping_mask.view(1,-1) # (m,c)
        out_sam_mask_cls = sam_mask_cls * (1 - category_overlapping_mask).view(1,-1) # (m,c)
        
        sam_mask_left = out_sam_mask_cls.max(dim=1).values>0.5 # m
        out_sam_mask_cls = out_sam_mask_cls[sam_mask_left] # m', c
        out_sam_mask_pred = sam_mask_pred[sam_mask_left] # m', h, w
        
        out_sam_mask_cls = out_sam_mask_cls * alpha
        in_sam_mask_cls = in_sam_mask_cls * beta

        in_semseg = torch.einsum("qc,qhw->chw", in_mask_cls, mask_pred)
        out_semseg = torch.einsum("qc,qhw->chw", out_mask_cls, mask_pred)
        if not out_sam_mask_cls.shape[0]==0:
            out_sam_semseg = torch.einsum("mc,mhw->chw", out_sam_mask_cls, out_sam_mask_pred)
        else:
            out_sam_semseg = torch.zeros_like(out_semseg)
        out_semseg = out_semseg + out_sam_semseg 
        semseg = in_semseg + out_semseg
        
        return semseg
    

    def instance_inference(self, mask_cls, mask_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]
        # if this is panoptic segmentation
        if self.panoptic_on:
            num_classes = len(self.test_metadata.stuff_classes)
        else:
            num_classes = len(self.test_metadata.thing_classes)
        labels = torch.arange(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 // 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 this is panoptic segmentation, we only keep the "thing" classes
        if self.panoptic_on:
            keep = torch.zeros_like(scores_per_image).bool()
            for i, lab in enumerate(labels_per_image):
                keep[i] = lab in self.test_metadata.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]

        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)
        # result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()

        # 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