# Copyright (c) Facebook, Inc. and its affiliates. from typing import Tuple import torch from torch import nn from torch.nn import functional as F import pdb import numpy as np import cv2 import os 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, BitMasks from detectron2.utils.memory import retry_if_cuda_oom from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES from .modeling.criterion import SetCriterion from .modeling.matcher import HungarianMatcher from .modeling.criterion_view import ViewSetCriterion from .modeling.matcher_view import ViewHungarianMatcher import pdb import copy @META_ARCH_REGISTRY.register() class CropFormer(nn.Module): """ Main class for mask classification semantic segmentation architectures. """ @configurable def __init__( self, *, cfg, backbone: Backbone, sem_seg_head: nn.Module, criterion_2d: nn.Module, criterion_3d: nn.Module, num_queries: int, object_mask_threshold: float, overlap_threshold: float, 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, ): """ 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.cfg = cfg self.backbone = backbone self.sem_seg_head = sem_seg_head self.criterion_2d = criterion_2d self.criterion_3d = criterion_3d ## colors self.colors = [info["color"] for info in COCO_CATEGORIES] 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) ## colors self.colors = [info["color"] for info in COCO_CATEGORIES] # 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 @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_2d = HungarianMatcher( cost_class=class_weight, cost_mask=mask_weight, cost_dice=dice_weight, num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS, ) matcher_3d = ViewHungarianMatcher( 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_2d = SetCriterion( sem_seg_head.num_classes, matcher=matcher_2d, 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, ) criterion_3d = ViewSetCriterion( sem_seg_head.num_classes, matcher=matcher_3d, 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, ) return { "cfg": cfg, "backbone": backbone, "sem_seg_head": sem_seg_head, "criterion_2d": criterion_2d, "criterion_3d": criterion_3d, "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, "metadata": MetadataCatalog.get(cfg.DATASETS.TRAIN[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, } @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". """ ## make new images batched_inputs_new = [] for batched_input in batched_inputs: ori_infos = {"height": batched_input["height"], "width": batched_input["width"], "image": batched_input["image"], # "file_name": batched_input["file_name"], } if "instances" in batched_input.keys(): ori_instances = batched_input["instances"] ori_instances.original_indices = torch.arange(0, len(ori_instances)).long() ori_infos["instances"] = ori_instances batched_inputs_new.append(ori_infos) ## cropped patches # pdb.set_trace() crop_region = batched_input["crop_region"] crop_images = batched_input["image_crop"] crop_o_width = int(crop_region[0][2]-crop_region[0][0]) crop_o_height = int(crop_region[0][3]-crop_region[0][1]) if "instances_crop" in batched_input.keys(): crop_instances = batched_input["instances_crop"] else: crop_instances = None for crop_index, crop_image in enumerate(crop_images): crop_infos = {"height": crop_o_height, "width": crop_o_width, "image": crop_image} if not crop_instances == None: crop_instance = crop_instances[crop_index] crop_instance.original_indices = torch.arange(0, len(crop_instance)).long() crop_infos["instances"] = crop_instance batched_inputs_new.append(crop_infos) images = [x["image"].to(self.device) for x in batched_inputs_new] ## +1 means num_views = self.cfg.ENTITY.CROP_SAMPLE_NUM_TRAIN+1 if self.training else self.cfg.ENTITY.CROP_SAMPLE_NUM_TEST+1 for i in range(len(images)): if i%num_views==0: continue _, c_h, c_w = images[i].shape if "instances" in batched_inputs_new[i].keys(): batched_inputs_new[i]["instances"]._image_size = (c_h, c_w) 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) outputs_2d, outputs_3d = self.sem_seg_head(features) if self.training: if self.cfg.ENTITY.ENABLE: for i in range(len(batched_inputs_new)): batched_inputs_new[i]["instances"].gt_classes[:] = 0 if "instances" in batched_inputs[0]: gt_instances = [x["instances"].to(self.device) for x in batched_inputs_new] targets_2d = self.prepare_targets_2d(copy.deepcopy(gt_instances), copy.deepcopy(images)) targets_3d = self.prepare_targets_3d(copy.deepcopy(gt_instances), copy.deepcopy(images), num_views) else: targets = None # bipartite matching-based loss losses = {} losses_2d = self.criterion_2d(outputs_2d, targets_2d) losses_3d = self.criterion_3d(outputs_3d, targets_3d) for k in list(losses_2d.keys()): if k in self.criterion_2d.weight_dict: losses[k+"_2d"] = losses_2d[k] * self.criterion_2d.weight_dict[k] * 0.5 else: # remove this loss if not specified in `weight_dict` losses_2d.pop(k) for k in list(losses_3d.keys()): if k in self.criterion_3d.weight_dict: losses[k+"_3d"] = losses_3d[k] * self.criterion_3d.weight_dict[k] else: # remove this loss if not specified in `weight_dict` losses_3d.pop(k) return losses else: mask_cls_results_3d = outputs_3d["pred_logits"][0] ## 100,2 mask_pred_results_3d = outputs_3d["pred_masks"][0] ## 100,5,200, 304 mask_cls_results_2d = outputs_2d["pred_logits"] mask_pred_results_2d = outputs_2d["pred_masks"] # upsample masks mask_pred_results_3d = retry_if_cuda_oom(F.interpolate)( mask_pred_results_3d, size=(images.tensor.shape[-2], images.tensor.shape[-1]), mode="bilinear", align_corners=False, ) mask_pred_results_2d = F.interpolate( mask_pred_results_2d, size=(images.tensor.shape[-2], images.tensor.shape[-1]), mode="bilinear", align_corners=False, ) del outputs_2d, outputs_3d crop_regions = batched_input["crop_region"][:num_views-1] processed_results = retry_if_cuda_oom(self.inference_whole_views)( mask_cls_results_3d, mask_pred_results_3d, mask_cls_results_2d, mask_pred_results_2d, batched_inputs_new, images.image_sizes, crop_regions) # processed_results = retry_if_cuda_oom(self.instance_inference_nonoverlap)( # mask_cls_results_2d[0], # mask_pred_results_2d[0], # batched_inputs_new[0], # images.image_sizes[0]) return [{"instances": processed_results}] def prepare_targets_2d(self, targets, images): h_pad, w_pad = images.tensor.shape[-2:] new_targets = [] for targets_per_image in targets: gt_masks = targets_per_image.gt_masks.tensor gt_valid = targets_per_image.gt_boxes_valid 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 valid_index = torch.nonzero(gt_valid).flatten() new_targets.append( { "labels": targets_per_image.gt_classes[valid_index], "masks": padded_masks[valid_index], } ) return new_targets def prepare_targets_3d(self, targets_ori, images, num_views): T = num_views B = int(len(targets_ori) / T) h_pad, w_pad = images.tensor.shape[-2:] ## reshape to new targets new_targets = [] for count, target in enumerate(targets_ori): b_index, t_index = int(count // T), int(count % T) if t_index == 0: new_targets.append([target]) else: new_targets[b_index].append(target) gt_instances = [] for count, targets in enumerate(new_targets): _num_instance = len(targets[0]) mask_shape = [_num_instance, T, h_pad, w_pad] gt_masks_per_view = torch.zeros(mask_shape, dtype=torch.bool, device=self.device) for v_i, targets_per_view in enumerate(targets): assert torch.all(targets[0].original_indices == targets_per_view.original_indices) gt_ids_per_view = [] gt_ids_per_valid = [] gt_ids_categories = [] ## view first, then entities for v_i, targets_per_view in enumerate(targets): targets_per_view = targets_per_view.to(self.device) h, w = targets_per_view.image_size for i_i, (instance_mask, instance_valid) in enumerate(zip(targets_per_view.gt_masks.tensor, targets_per_view.gt_boxes_valid)): if instance_valid == 1: gt_masks_per_view[i_i, v_i, :h, :w] = instance_mask gt_ids_per_valid.append(targets_per_view.gt_boxes_valid[None,:]) gt_ids_per_view.append(targets_per_view.original_indices[None,:]) gt_ids_categories.append(targets_per_view.gt_classes[None, :]) ## (num_instances, num_views) gt_ids_per_valid = torch.cat(gt_ids_per_valid, dim=0).permute((1,0)) gt_ids_per_view = torch.cat(gt_ids_per_view, dim=0).permute((1,0)) gt_ids_categories = torch.cat(gt_ids_categories, dim=0).permute((1,0)) gt_ids_per_view[gt_ids_per_valid == 0] = -1 valid_idx = (gt_ids_per_view != 1).any(dim=-1) ## categoreis gt_classes_per_group = gt_ids_categories[:,0] ## N gt_ids_per_group = gt_ids_per_view ## N, num_views gt_masks_per_group = gt_masks_per_view.float() ## N, num_views, H, W ## gt_instances.append({"labels": gt_classes_per_group, "ids": gt_ids_per_group, "masks": gt_masks_per_group}) return gt_instances 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() 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 = {} for k in range(cur_classes.shape[0]): pred_class = cur_classes[k].item() isthing = pred_class in self.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: 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_nonoverlap(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] 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] ###### ranks pred_masks = (mask_pred>0).float() pred_masks_logits = mask_pred.sigmoid() pred_scores = scores_per_image _, m_H, m_W = pred_masks.shape mask_id = torch.zeros((m_H, m_W), dtype=torch.int).to(pred_masks.device) sorted_scores, ranks = torch.sort(pred_scores) ranks = ranks + 1 for index in ranks: mask_id[(pred_masks[index-1]==1)] = int(index) # re-generate mask new_scores = [] new_masks = [] new_masks_logits = [] entity_nums = len(ranks) for ii in range(entity_nums): index = int(ranks[entity_nums-ii-1]) score = sorted_scores[entity_nums-ii-1] new_scores.append(score) new_masks.append((mask_id==index).float()) new_masks_logits.append(pred_masks_logits[index-1]) new_scores = torch.stack(new_scores) new_masks = torch.stack(new_masks) new_masks_logits = torch.stack(new_masks_logits) result = Instances(image_size) # mask (before sigmoid) result.pred_masks = new_masks result.pred_boxes = Boxes(torch.zeros(new_masks.size(0), 4)) # Uncomment the following to get boxes from masks (this is slow) # calculate average mask prob mask_scores_per_image = (new_masks_logits.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6) result.scores = new_scores * mask_scores_per_image result.pred_classes = labels_per_image return result 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] 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 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.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) # pdb.set_trace() result.scores = scores_per_image * mask_scores_per_image result.pred_classes = labels_per_image return result def inference_whole_views(self, pred_cls, pred_masks, pred_cls_2d, pred_masks_2d, batched_inputs, image_sizes, crop_regions): ## pred_masks: [100, 5, 800, 1216] ## pred_masks_2d: [5, 100, 800, 1216] scores = F.softmax(pred_cls, dim=-1)[:,:-1] # 100,1 scores_2d = F.softmax(pred_cls_2d, dim=-1)[:, :, :-1] # 5, 100, 1 # scores = (scores+scores_2d[0])/2 labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1) ### keep all the indices scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False) labels_per_image = labels[topk_indices] # topk_indices = topk_indices // self.sem_seg_head.num_classes topk_indices = torch.div(topk_indices, self.sem_seg_head.num_classes, rounding_mode="trunc") pred_masks = pred_masks[topk_indices] pred_masks = pred_masks.permute((1,0,2,3)) new_pred_masks = [] for view_index, (pred_masks_per_view, batched_input_per_view, image_size_per_view) in enumerate(zip(pred_masks, batched_inputs, image_sizes)): O_H = batched_input_per_view["height"] O_W = batched_input_per_view["width"] SO_H, SO_W = image_size_per_view pred_masks_per_view = pred_masks_per_view[..., : SO_H, :SO_W] pred_masks_per_view = F.interpolate(pred_masks_per_view[None], size=(O_H, O_W), mode="bilinear", align_corners=False) new_pred_masks.append(pred_masks_per_view[0].sigmoid()) ## fuse the masks full_image_masks = new_pred_masks[0] ## fuse crop image fused_image_masks = torch.zeros_like(full_image_masks).float() fused_image_masks_valid = torch.zeros_like(full_image_masks).float() + 1e-16 for crop_region_per_view, pred_masks_per_view in zip(crop_regions, new_pred_masks[1:]): x0, y0, x1, y1 = crop_region_per_view fused_image_masks[..., y0:y1, x0:x1] += pred_masks_per_view fused_image_masks_valid[..., y0:y1, x0:x1] += 1 # add original masks fused_image_masks += full_image_masks fused_image_masks_valid += 1 ## average fuse_image_masks = fused_image_masks / fused_image_masks_valid ###### change to the single image, begin to non_overlap_supression ## ranks pred_masks_logits = fuse_image_masks pred_masks = (fuse_image_masks>0.5).float() pred_scores = scores_per_image _, m_H, m_W = pred_masks.shape ## for visualization mask_id = torch.zeros((m_H, m_W), dtype=torch.int).to(pred_masks.device) # mask_id_colors = np.zeros((m_H, m_W, 3), dtype=np.uint8) # pred_masks_np = pred_masks.cpu().numpy() sorted_scores, ranks = torch.sort(pred_scores) ranks = ranks + 1 for index in ranks: mask_id[(pred_masks[index-1]==1)] = int(index) # mask_id_colors[(pred_masks_np[index-1]==1)] = self.colors[index] # base_path = "/group/20018/gavinqi/vis_entityv2_release_debug" # pdb.set_trace() # file_name = batched_inputs[0]["file_name"] # split_index, img_name = file_name.split("/")[-2:] # save_name = img_name.split(".")[0]+".png" # if not os.path.exists(os.path.join(base_path, save_name)): # cv2.imwrite(os.path.join(base_path, save_name), mask_id_colors) # re-generate mask new_scores = [] new_masks = [] new_masks_logits = [] entity_nums = len(ranks) for ii in range(entity_nums): index = int(ranks[entity_nums-ii-1]) score = sorted_scores[entity_nums-ii-1] new_scores.append(score) new_masks.append((mask_id==index).float()) new_masks_logits.append(pred_masks_logits[index-1]) new_scores = torch.stack(new_scores) new_masks = torch.stack(new_masks) new_masks_logits = torch.stack(new_masks_logits) # make result image_size = (batched_inputs[0]["height"], batched_inputs[0]["width"]) result = Instances(image_size) # mask (before sigmoid) result.pred_masks = new_masks result.pred_boxes = Boxes(torch.zeros(new_masks.size(0), 4)) # Uncomment the following to get boxes from masks (this is slow) # calculate average mask prob mask_scores_per_image = (new_masks_logits.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6) result.scores = new_scores * mask_scores_per_image result.pred_classes = labels_per_image return result