""" This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”). All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates. Reference: https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/maskformer_model.py """ from typing import Tuple import os import numpy as np import torch from torch import nn from torch.nn import functional as F from torchvision import transforms as T from detectron2.config import configurable from detectron2.data import MetadataCatalog from detectron2.modeling import META_ARCH_REGISTRY, build_backbone from detectron2.modeling.backbone import Backbone from detectron2.modeling.postprocessing import sem_seg_postprocess from detectron2.structures import Boxes, ImageList, Instances, BitMasks from detectron2.utils.memory import retry_if_cuda_oom from .modeling.maft.content_dependent_transfer import ContentDependentTransfer from .modeling.meta_arch.mask_adapter_head import build_mask_adapter 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 MASK_Adapter(nn.Module): """ Main class for mask classification semantic segmentation architectures. """ @configurable def __init__( self, *, backbone: Backbone, mask_adapter: nn.Module, weight_dict, num_queries: int, object_mask_threshold: float, overlap_threshold: float, mask_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, train_maft : bool, num_output_maps: int, ): """ Args: backbone: a backbone module, must follow detectron2's backbone interface mask_adapter: mask-adapter extract semantic activation maps from masks weight_dict: dict contains weight for each 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.mask_adapter = mask_adapter self.weight_dict = weight_dict self.num_queries = num_queries self.overlap_threshold = overlap_threshold self.object_mask_threshold = object_mask_threshold self.mask_threshold = 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.void_embedding = nn.Embedding(1, backbone.dim_latent) self.train_dataname = None self.test_dataname = None self.train_num_templates = {} self.train_text_classifier = {} self.train_maft = train_maft self.num_output_maps = num_output_maps if self.train_maft: if '_base' in backbone.model_name.lower(): cdt_params = [640, 8] elif '_large' in backbone.model_name.lower(): cdt_params = [768, 8] self.cdt = ContentDependentTransfer(d_model = cdt_params[0], nhead = cdt_params[1], panoptic_on = panoptic_on) self.freeze_cdt() def freeze_cdt(self): for param in self.cdt.parameters(): param.requires_grad = False #https://github.com/bytedance/fc-clip/blob/2b0bbe213070d44da9182530fa2e826fef03f974/fcclip/fcclip.py#L139 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 #print("text for classification:", 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, dataname): if self.training: os.makedirs("text_embedding", exist_ok=True) out_path = f"./text_embedding/{dataname}_text_embedding.npy" if dataname in self.train_text_classifier: return self.train_text_classifier[dataname], self.train_num_templates[dataname] if dataname not in self.train_num_templates: _, self.train_num_templates[dataname], train_class_names = self.prepare_class_names_from_metadata( self.train_metadata[dataname], self.train_metadata[dataname] ) if os.path.exists(out_path): text_classifier = torch.from_numpy(np.load(out_path)).to(self.device) else: text_classifier = [] bs = 128 for idx in range(0, len(train_class_names), bs): text_classifier.append( self.backbone.get_text_classifier(train_class_names[idx:idx+bs], self.device).detach() ) text_classifier = torch.cat(text_classifier, dim=0) 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) np.save(out_path, text_classifier.cpu().numpy()) self.train_text_classifier[dataname] = text_classifier return self.train_text_classifier[dataname], self.train_num_templates[dataname] else: if self.test_dataname != dataname: self.category_overlapping_mask, self.test_num_templates, self.test_class_names = self.prepare_class_names_from_metadata( self.test_metadata[dataname], self.test_metadata[dataname] ) text_classifier = [] 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) 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 self.test_dataname = dataname return self.test_text_classifier, self.test_num_templates @classmethod def from_config(cls, cfg): backbone = build_backbone(cfg) mask_adapter = build_mask_adapter(cfg, cfg.MODEL.MASK_ADAPTER.NAME) # loss weights class_weight = cfg.MODEL.MASK_FORMER.CLASS_WEIGHT # building criterion weight_dict = {"loss_ce": class_weight} losses = ["labels"] train_metadata = {i: MetadataCatalog.get(i) for i in cfg.DATASETS.TRAIN} test_metadata = {i: MetadataCatalog.get(i) for i in cfg.DATASETS.TEST} return { "backbone": backbone, "mask_adapter": mask_adapter, "weight_dict": weight_dict, "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, "mask_threshold": cfg.MODEL.MASK_ADAPTER.MASK_THRESHOLD, "train_metadata": train_metadata,#MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), "test_metadata": 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 ), "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, "test_topk_per_image": cfg.TEST.DETECTIONS_PER_IMAGE, "train_maft": cfg.MODEL.MASK_ADAPTER.TRAIN_MAFT, "num_output_maps": cfg.MODEL.MASK_ADAPTER.NUM_OUTPUT_MAPS } @property def device(self): return self.pixel_mean.device 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": A Tensor that represents the per-pixel segmentation prediced by the head. The prediction has shape KxHxW that represents the logits of each class for each pixel. * "panoptic_seg": A tuple that represent panoptic output panoptic_seg (Tensor): of shape (height, width) where the values are ids for each segment. segments_info (list[dict]): Describe each segment in `panoptic_seg`. Each dict contains keys "id", "category_id", "isthing". """ if self.train_maft and self.training : dataname = "openvocab_coco_2017_train_stuff_sem_seg" else: dataname = batched_inputs[0]['dataname'] if self.training: dataname_2 = batched_inputs[1]['dataname'] assert dataname == dataname_2, f"expect batch img from same dataset, but different from {dataname} and {dataname_2}" 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) features = self.backbone(images.tensor) clip_feature = features['clip_vis_dense'] text_classifier, num_templates = self.get_text_classifier(dataname) text_classifier = torch.cat([text_classifier, F.normalize(self.void_embedding.weight, dim=-1)], dim=0) clip_vis_dense = self.visual_prediction_forward_convnext_2d(clip_feature) if self.train_maft: #https://github.com/jiaosiyu1999/MAFT-Plus/blob/fd12806df651d309883229de9503e40533f92689/maft/maft_plus.py#L352 #For maftp,it uses a wrong reshape operation to get clip_vis_dense. Since we don't finetune cdt, we follow them. img_feat = self.visual_prediction_forward_convnext(clip_feature) text_classifier = self.cdt(img_feat, text_classifier) clip_vis_dense = img_feat else: clip_vis_dense = self.visual_prediction_forward_convnext_2d(clip_feature) if self.training: # mask classification target if "instances" in batched_inputs[0]: gt_instances = [x["instances"].to(self.device) for x in batched_inputs] targets,masks,labels = self.prepare_targets(gt_instances, images) else: targets = None semantic_activation_maps = self.mask_adapter(clip_vis_dense, masks) maps_for_pooling = F.interpolate(semantic_activation_maps, size=clip_feature.shape[-2:], mode='bilinear', align_corners=False) if "convnext" in self.backbone.model_name.lower(): B, C = clip_feature.size(0),clip_feature.size(1) N = maps_for_pooling.size(1) num_instances = N // self.num_output_maps maps_for_pooling = F.softmax(F.logsigmoid(maps_for_pooling).view(B, N,-1), dim=-1) pooled_clip_feature = torch.bmm(maps_for_pooling, clip_feature.view(B, C, -1).permute(0, 2, 1)) pooled_clip_feature = self.backbone.visual_prediction_forward(pooled_clip_feature) pooled_clip_feature = (pooled_clip_feature.reshape(B,num_instances, self.num_output_maps, -1).mean(dim=-2).contiguous()) else: raise NotImplementedError mask_cls_results = get_classification_logits(pooled_clip_feature, text_classifier, self.backbone.clip_model.logit_scale, num_templates) losses = self.cross_entropy_loss(mask_cls_results, labels) for k in list(losses.keys()): if k in self.weight_dict: losses[k] *= self.weight_dict[k] else: # remove this loss if not specified in `weight_dict` losses.pop(k) return losses else: masks = [] classes = [] for input_per_image in batched_inputs: height = input_per_image.get("height") width = input_per_image.get("width") sem_seg = input_per_image["sem_seg"].to(self.device) total_masks,class_label = self.sem_seg_2_gt_masks(sem_seg, height, width) masks.append(total_masks) classes.append(class_label) masks = torch.stack(masks) classes = torch.stack(classes) outputs = self.mask_adapter(clip_vis_dense, masks) maps_for_pooling = F.interpolate(outputs, size=clip_vis_dense.shape[-2:], mode='bilinear', align_corners=False) if "convnext" in self.backbone.model_name.lower(): B,C = clip_feature.size(0),clip_feature.size(1) N = maps_for_pooling.size(1) num_instances = N // self.num_output_maps maps_for_pooling = F.softmax(F.logsigmoid(maps_for_pooling).view(B, N,-1), dim=-1) pooled_clip_feature = torch.bmm(maps_for_pooling, clip_feature.view(B, C, -1).permute(0, 2, 1)) pooled_clip_feature = self.backbone.visual_prediction_forward(pooled_clip_feature) pooled_clip_feature = (pooled_clip_feature.reshape(B,num_instances, self.num_output_maps, -1).mean(dim=-2).contiguous()) else: raise NotImplementedError mask_cls_results = get_classification_logits(pooled_clip_feature, text_classifier, self.backbone.clip_model.logit_scale, num_templates) mask_cls_results = mask_cls_results.softmax(-1) #upsample masks mask_pred_results = F.interpolate( masks, size=(images.tensor.shape[-2], images.tensor.shape[-1]), mode="bilinear", align_corners=False, ) 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: 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) mask_pred_result = mask_pred_result.squeeze(1) # semantic segmentation inference if self.semantic_on: r = retry_if_cuda_oom(self.semantic_inference)(mask_cls_result, mask_pred_result) 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: 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 sem_seg_2_gt_masks(self, sem_seg, height, width): classes = torch.unique(sem_seg,sorted=False,return_inverse=False,return_counts=False) gt_labels = classes[classes != 255] masks = [sem_seg == class_id for class_id in gt_labels] if len(masks) == 0: gt_masks = torch.zeros((0, sem_seg.shape[-2], sem_seg.shape[-1])).to(sem_seg) else: gt_masks = torch.stack(masks).squeeze(1) num_masks = gt_masks.shape[0] total_masks = torch.zeros((num_masks, gt_masks.shape[1], gt_masks.shape[2]), dtype=gt_masks.dtype, device=gt_masks.device) labels = torch.zeros((num_masks), device=gt_masks.device) total_masks[:num_masks] = gt_masks[:num_masks] labels[:num_masks] = gt_labels[:num_masks] return total_masks.float(), labels def visual_prediction_forward_convnext(self, x): batch, channel, h, w = x.shape x = x.reshape(batch*h*w, channel).unsqueeze(-1).unsqueeze(-1) # fake 2D input x = self.backbone.clip_model.visual.trunk.head(x) x = self.backbone.clip_model.visual.head(x) return x.reshape(batch, h, w, x.shape[-1]).permute(0,3,1,2) def visual_prediction_forward_convnext_2d(self, x): clip_vis_dense = self.backbone.clip_model.visual.trunk.head.norm(x) clip_vis_dense = self.backbone.clip_model.visual.trunk.head.drop(clip_vis_dense.permute(0, 2, 3, 1)) clip_vis_dense = self.backbone.clip_model.visual.head(clip_vis_dense).permute(0, 3, 1, 2) return clip_vis_dense def cross_entropy_loss(self, mask_cls_results, labels): if torch.all(labels == -1): loss_ce = mask_cls_results.sum() * 0.0 else: loss_ce = F.cross_entropy(mask_cls_results.transpose(1, 2), labels.to(torch.int64), ignore_index=-1) #remove celoss weight because of multiple datasets training losses = {"loss_ce": loss_ce} return losses def prepare_targets(self, targets, images): h_pad, w_pad = images.tensor.shape[-2:] new_targets = [] masks_list = [] labels_list = [] num_masks = 32 min_mask_area = 0 for targets_per_image in targets: gt_masks = targets_per_image.gt_masks if isinstance(gt_masks, BitMasks): gt_masks = gt_masks.tensor valid_mask_indices = [i for i, mask in enumerate(gt_masks) if mask.sum() > min_mask_area] if len(valid_mask_indices) > 0: valid_gt_masks = gt_masks[valid_mask_indices] valid_gt_classes = targets_per_image.gt_classes[valid_mask_indices] padded_masks = torch.zeros((valid_gt_masks.shape[0], h_pad, w_pad), dtype=valid_gt_masks.dtype, device=valid_gt_masks.device) padded_masks[:, : valid_gt_masks.shape[1], : valid_gt_masks.shape[2]] = valid_gt_masks new_targets.append( { "labels": valid_gt_classes, "masks": padded_masks, } ) total_masks = torch.zeros((num_masks, h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device) selected_labels = torch.zeros((num_masks), device=gt_masks.device) if valid_gt_masks.shape[0] > num_masks: selected_indices = torch.randperm(valid_gt_masks.shape[0])[:num_masks] for idx, mask_idx in enumerate(selected_indices): total_masks[idx, :valid_gt_masks[mask_idx].shape[0], :valid_gt_masks[mask_idx].shape[1]] = valid_gt_masks[mask_idx] selected_labels[idx] = valid_gt_classes[mask_idx] else: for idx in range(valid_gt_masks.shape[0]): total_masks[idx, :valid_gt_masks[idx].shape[0], :valid_gt_masks[idx].shape[1]] = valid_gt_masks[idx] selected_labels[idx] = valid_gt_classes[idx] for idx in range(valid_gt_masks.shape[0], num_masks): total_masks[idx] = torch.zeros((h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device) selected_labels[idx] = -1 else: total_masks = torch.zeros((num_masks, h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device) selected_labels = torch.zeros((num_masks), device=gt_masks.device) selected_labels.fill_(-1) padded_masks = torch.zeros((0, h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device) valid_gt_classes = torch.zeros((0), device=gt_masks.device) new_targets.append( { "labels": valid_gt_classes, "masks": padded_masks, } ) masks_list.append(total_masks) labels_list.append(selected_labels) masks = torch.stack(masks_list, dim=0) labels = torch.stack(labels_list, dim=0) labels = labels.long() return new_targets, masks, labels def semantic_inference(self, mask_cls, mask_pred): mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1] if mask_pred.dim() == 4: mask_pred = mask_pred.squeeze(dim=0) #mask_pred = mask_pred.sigmoid() #remove because of gt masks 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) num_classes = len(self.test_metadata[self.test_dataname].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]): pred_class = cur_classes[k].item() isthing = pred_class in self.test_metadata[self.test_dataname].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): # 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] #[250,150] scores = mask_cls[:, :-1].sigmoid() # if this is panoptic segmentation if self.panoptic_on: num_classes = len(self.test_metadata[self.test_dataname].stuff_classes) else: num_classes = len(self.test_metadata[self.test_dataname].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[self.test_dataname].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 > self.mask_threshold).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.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 class MaskPooling(nn.Module): def __init__( self,mask_threshold ): super().__init__() self.mask_threshold = mask_threshold def forward(self, x, mask): """ Args: x: [B, C, H, W] mask: [B, Q, H, W] """ if not x.shape[-2:] == mask.shape[-2:]: # reshape mask to x mask = F.interpolate(mask, size=x.shape[-2:], mode='bilinear', align_corners=False) with torch.no_grad(): mask = mask.detach() binary_mask = (mask > self.mask_threshold).to(mask.dtype) mask = binary_mask * mask denorm = mask.sum(dim=(-1, -2), keepdim=True) + 1e-8 mask_pooled_x = torch.einsum( "bchw,bqhw->bqc", x, mask / denorm, ) return mask_pooled_x def get_classification_logits(x, text_classifier, logit_scale, num_templates=None): # x in shape of [B, *, C] # text_classifier in shape of [num_classes, C] # logit_scale is a learnable scalar https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/model.py#L201 # return: [B, *, num_classes] x = F.normalize(x, dim=-1) logit_scale = torch.clamp(logit_scale.exp(), max=100) if len(text_classifier.shape) == 2: pred_logits = logit_scale * x @ text_classifier.T # B, *, N + 1 else: pred_logits = logit_scale * x @ text_classifier.permute(0,2,1) # B, *, N + 1 # max ensembel as in OpenSeg/ODISE if pred_logits.shape[2] != 1204 and pred_logits.shape[2] != 366: final_pred_logits = [] cur_idx = 0 for num_t in num_templates: final_pred_logits.append(pred_logits[:, :, cur_idx: cur_idx + num_t].max(-1).values) cur_idx += num_t final_pred_logits.append(pred_logits[:, :, -1]) # the last classifier is for void final_pred_logits = torch.stack(final_pred_logits, dim=-1) else: final_pred_logits = pred_logits return final_pred_logits