from typing import List, Tuple import os import torch import torch.nn as nn import torch.nn.functional as F import torchvision from collections import defaultdict # modeing from transformers import AutoTokenizer from .evf_sam2 import EvfSam2Model from .criterion import SetCriterion from .matcher import HungarianMatcher from .segment_anything_2.sam2.modeling.sam2_utils import MLP from detectron2.config import configurable from detectron2.modeling import META_ARCH_REGISTRY, build_backbone, build_sem_seg_head from detectron2.structures import Boxes, ImageList, Instances, BitMasks from detectron2.utils.memory import retry_if_cuda_oom from detectron2.data import MetadataCatalog import logging @META_ARCH_REGISTRY.register() class OpenWorldSAM2(nn.Module): @configurable def __init__( self, *, evf_sam2: EvfSam2Model, tokenizer: AutoTokenizer, visual_model: nn.Module, mm_extractor: nn.Module, text_hidden_fcs: nn.ModuleList, query_dim: int, num_tokens: int, positional_tokens: nn.Parameter, criterion: nn.Module, pixel_mean: Tuple[float], pixel_std: Tuple[float], dtype: torch.dtype, test_topk_per_image: int, top_k_on: bool, nms_on: bool, nms_threshold: float, iou_threshold: float, semantic_on: bool, instance_on: bool, panoptic_on: bool, use_visual_tokens: bool = True, use_cross_attention: bool = False, cross_attention_layers: int = 3, # Added parameter for number of layers two_stage_inference: bool = False, # Add new parameter here refer_on: bool = False, # Add refer_on parameter metadata: MetadataCatalog = None, ): super(OpenWorldSAM2, self).__init__() self.evf_sam2 = evf_sam2 self.tokenizer = tokenizer self.visual_model = visual_model self.mm_extractor = mm_extractor self.text_hidden_fcs = text_hidden_fcs self.query_dim = query_dim # query embedding dimension self.num_tokens = num_tokens self.criterion = criterion self.positional_tokens = positional_tokens self.use_visual_tokens = use_visual_tokens self.use_cross_attention = use_cross_attention self.metadata = metadata self.two_stage_inference = two_stage_inference # Store the new parameter self.refer_on = refer_on # Store refer_on parameter # Add cross-attention transformer if enabled if self.use_cross_attention: self.cross_attention_transformer = CrossAttentionTransformer( embedding_dim=256, num_heads=8, mlp_dim=query_dim * 4, num_layers=cross_attention_layers, # Use the new parameter dropout=0.1 ) 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) self.dtype = dtype # additional args self.semantic_on = semantic_on self.instance_on = instance_on self.panoptic_on = panoptic_on self.top_k_on = top_k_on self.nms_on = nms_on self.test_topk_per_image = test_topk_per_image self.nms_threshold = nms_threshold self.iou_threshold = iou_threshold self._bb_feat_sizes = [ (256, 256), (128, 128), (64, 64), ] @classmethod def from_config(cls, cfg): # EVF-SAM config & model evf_config = cfg.MODEL.OpenWorldSAM2.EVF_CONFIG torch_dtype = torch.float32 kwargs = {"torch_dtype": torch_dtype} # tokenizer tokenizer_config = cfg.MODEL.OpenWorldSAM2.TOKENIZER_CONFIG tokenizer = AutoTokenizer.from_pretrained(tokenizer_config, padding_side="right", use_fast=False) # EVF-SAM2 model evf_sam2 = EvfSam2Model.from_pretrained(evf_config, low_cpu_mem_usage=True, **kwargs) evf_sam2.config.eos_token_id = tokenizer.eos_token_id evf_sam2.config.bos_token_id = tokenizer.bos_token_id evf_sam2.config.pad_token_id = tokenizer.pad_token_id # SAM2 visual model visual_model = evf_sam2.visual_model print("Loading SAM2 model from {}...".format(cfg.MODEL.OpenWorldSAM2.VISION_PRETRAINED)) visual_model.load_state_dict(torch.load(cfg.MODEL.OpenWorldSAM2.VISION_PRETRAINED)["model"], strict=False) for param in visual_model.parameters(): param.requires_grad = False # BEiT-3 model mm_extractor = evf_sam2.mm_extractor if cfg.MODEL.OpenWorldSAM2.TRAIN_VLM: for param in mm_extractor.parameters(): param.requires_grad = True else: for param in mm_extractor.parameters(): param.requires_grad = False # Projection Layer query_dim = cfg.MODEL.OpenWorldSAM2.QUERY_DIM in_dim = evf_sam2.config.hidden_size text_fc = [ nn.Linear(in_dim, in_dim), nn.ReLU(), nn.Linear(in_dim, query_dim) ] text_hidden_fcs = nn.ModuleList([nn.Sequential(*text_fc)]) text_hidden_fcs.train() for param in text_hidden_fcs.parameters(): param.requires_grad = True # OpenWorldSAM2 config num_tokens = cfg.MODEL.OpenWorldSAM2.NUM_OBJECT_QUERIES positional_tokens = nn.Parameter(torch.randn(num_tokens, query_dim)) positional_tokens.requires_grad = True # Loss parameters: no_object_weight = cfg.MODEL.OpenWorldSAM2.NO_OBJECT_WEIGHT dice_weight = cfg.MODEL.OpenWorldSAM2.DICE_WEIGHT mask_weight = cfg.MODEL.OpenWorldSAM2.MASK_WEIGHT objectness_weight = cfg.MODEL.OpenWorldSAM2.OBJECTNESS_WEIGHT # Get use_cross_attention from config use_cross_attention = getattr(cfg.MODEL.OpenWorldSAM2, "USE_CROSS_ATTENTION", False) # Get two_stage_inference from config with default=False two_stage_inference = getattr(cfg.MODEL.OpenWorldSAM2.TEST, "TWO_STAGE_INFERENCE", False) # Get refer_on from config with default=False refer_on = getattr(cfg.MODEL.OpenWorldSAM2.TEST, "REFER_ON", False) # building criterion matcher = HungarianMatcher( cost_class=objectness_weight, cost_mask=mask_weight, cost_dice=dice_weight, ) weight_dict = {"loss_ce": objectness_weight, "loss_mask": mask_weight, "loss_dice": dice_weight} losses = ["labels", "masks"] criterion = SetCriterion( num_classes=1, # omitting the special no-object category, 1 class to indicate object or no object matcher=matcher, weight_dict=weight_dict, eos_coef=no_object_weight, losses=losses, ) return { "evf_sam2": evf_sam2, "tokenizer": tokenizer, "visual_model": visual_model, "mm_extractor": mm_extractor, "text_hidden_fcs": text_hidden_fcs, "query_dim": query_dim, "num_tokens": num_tokens, "positional_tokens": positional_tokens, "criterion": criterion, "pixel_mean": cfg.MODEL.PIXEL_MEAN, "pixel_std": cfg.MODEL.PIXEL_STD, "dtype": torch_dtype, # inference "semantic_on": cfg.MODEL.OpenWorldSAM2.TEST.SEMANTIC_ON, "instance_on": cfg.MODEL.OpenWorldSAM2.TEST.INSTANCE_ON, "panoptic_on": cfg.MODEL.OpenWorldSAM2.TEST.PANOPTIC_ON, "top_k_on": cfg.MODEL.OpenWorldSAM2.TEST.TOP_K_ON, "nms_on": cfg.MODEL.OpenWorldSAM2.TEST.NMS_ON, "test_topk_per_image": cfg.MODEL.OpenWorldSAM2.TEST.DETECTIONS_PER_IMAGE, "nms_threshold": cfg.MODEL.OpenWorldSAM2.TEST.NMS_THRESHOLD, "iou_threshold": cfg.MODEL.OpenWorldSAM2.TEST.IOU_THRESHOLD, "use_visual_tokens": cfg.MODEL.OpenWorldSAM2.USE_VISUAL_TOKENS, "use_cross_attention": use_cross_attention, "cross_attention_layers": cfg.MODEL.OpenWorldSAM2.CROSS_ATTENTION_LAYERS, "two_stage_inference": two_stage_inference, # Add the new parameter here "refer_on": refer_on, # Add refer_on from config "metadata": MetadataCatalog.get(cfg['DATASETS']['TRAIN'][0]) , } def print_trainable_parameters(self): """ Prints the names and number of trainable parameters in the model. """ logger = logging.getLogger("detectron2") total_params = 0 trainable_params = 0 logger.info(f"{'Parameter Name':<40}{'Trainable':<10}{'Shape':<20}{'Num Params':<15}") logger.info("=" * 85) for name, param in self.named_parameters(): num_params = param.numel() total_params += num_params if param.requires_grad: trainable_params += num_params trainable_status = "Yes" logger.info(f"{name:<40}{trainable_status:<10}{str(list(param.shape)):<20}{num_params:<15}") else: trainable_status = "No" logger.info("=" * 85) logger.info(f"Total parameters: {total_params}") logger.info(f"Trainable parameters: {trainable_params}") logger.info(f"Non-trainable parameters: {total_params - trainable_params}") logger.info("=" * 85) logger.info(f"use_cross_attention: {self.use_cross_attention}") logger.info(f"use_visual_tokens: {self.use_visual_tokens}") logger.info(f"two_stage_inference: {self.two_stage_inference}") @property def device(self): return self.pixel_mean.device def tokenize_prompts(self, prompts: List): input_ids = [ self.tokenizer(prompt, return_tensors="pt").input_ids[0] for prompt in prompts ] input_ids = torch.nn.utils.rnn.pad_sequence( input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id ) attention_masks = input_ids.ne(self.tokenizer.pad_token_id) truncate_len = self.tokenizer.model_max_length if input_ids.shape[1] > truncate_len: input_ids = input_ids[:, :truncate_len] attention_masks = attention_masks[:, :truncate_len] return input_ids.to(self.device), attention_masks.to(self.device) def forward( self, batched_inputs, return_intermediate=False ): """ 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. * prompts: a list of prompts for the corresponding image * unique_categories: unique IDs for the corresponding prompt Returns: dict[str, Tensor]: """ ######################## input pre-processing ####################### images = [x["image"].to(dtype=self.dtype, device=self.device) for x in batched_inputs] original_size_list = [(x["height"], x["width"]) for x in batched_inputs] images_evf = [x["evf_image"].to(dtype=self.dtype, device=self.device) for x in batched_inputs] # Convert to tensors images = ImageList.from_tensors(images, 1024).tensor images_evf = ImageList.from_tensors(images_evf, 224).tensor # Calculate offsets for prompts per image offset = [0] all_prompts = [] # Process each image and build prompts list for x in batched_inputs: prompts = x["prompt"] all_prompts.extend(prompts) offset.append(offset[-1] + len(prompts)) input_ids, attention_masks = self.tokenize_prompts(all_prompts) batch_size = len(batched_inputs) assert batch_size == len(offset) - 1 ############################## forward ############################# backbone_out = self.visual_model.forward_image(images) # dict_keys(['vision_features', 'vision_pos_enc', 'backbone_fpn']) _, image_embeddings, _, _ = self.visual_model._prepare_backbone_features(backbone_out) # Expand images_evf according to number of prompts per image if self.use_visual_tokens: images_evf_list = [] for i in range(len(offset) - 1): start_i, end_i = offset[i], offset[i + 1] images_evf_i = ( images_evf[i] .unsqueeze(0) .expand(end_i - start_i, -1, -1, -1) .contiguous() ) images_evf_list.append(images_evf_i) images_evf = torch.cat(images_evf_list, dim=0) # Process through BEIT-3 output = self.mm_extractor.beit3( visual_tokens=images_evf, textual_tokens=input_ids, text_padding_position=~attention_masks, ) else: # When not using visual tokens, we'll pass None output = self.mm_extractor.beit3( visual_tokens=None, textual_tokens=input_ids, text_padding_position=~attention_masks, ) feat = output["encoder_out"][:, :1, ...] feat = self.text_hidden_fcs[0](feat) # Split features back according to images """ Within a single image of the training dataset, there are several (usually more than 1) referring expressions corresponding to different parts of the image. For example we use batch 2 to train the code, and the first image has 3 referring expressions and the secode image has 2 referring expressions, the offset would be [0, 3, 5]. The torch.split would split the multi-modal extracted feat of length 5 to a list, where each item of the list corresponds to each image in batch. """ feat = torch.split(feat, [offset[i + 1] - offset[i] for i in range(len(offset) - 1)]) # print(f"Split features length: {len(feat)}, First feature shape: {feat[0].shape}") # Process image features image_embeddings = [_.to(images.dtype) for _ in image_embeddings] if self.visual_model.directly_add_no_mem_embed: image_embeddings[-1] = image_embeddings[-1] + self.visual_model.no_mem_embed feats = [ feat.permute(1, 2, 0).view(batch_size, -1, *feat_size) for feat, feat_size in zip(image_embeddings[::-1], self._bb_feat_sizes[::-1]) ][::-1] _features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]} if self.training: # Initialize lists to store all predictions and losses all_losses = defaultdict(list) if not self.training: processed_results = [] # Process each image batch for img_idx in range(batch_size): img_feat = feat[img_idx] # Get features for all prompts of this image # Prepare all feat_with_tokens for this image's prompts batch_feat_with_tokens = [] for prompt_idx, prompt_feat in enumerate(img_feat): # Repeat feature along token dimension and add positional embeddings feat_repeated = prompt_feat.expand(self.num_tokens, -1, -1) feat_with_tokens = feat_repeated + self.positional_tokens.unsqueeze(1) batch_feat_with_tokens.append(feat_with_tokens) # Concatenate all prompts for this image batch_feat_with_tokens = torch.cat(batch_feat_with_tokens, dim=0) # Apply cross-attention if enabled if self.use_cross_attention: # Prepare image embeddings for cross-attention img_embed = _features["image_embed"][img_idx] # [C, H, W] img_embed = img_embed.flatten(1).transpose(0, 1) # [H*W, C] # Add a batch dimension to img_embed to make it 3D [1, H*W, C] img_embed = img_embed.unsqueeze(0) # Apply cross-attention original_batch_feat_with_tokens = batch_feat_with_tokens # Reshape batch_feat_with_tokens to be 3D [batch_size, num_tokens, embedding_dim] # The current shape is likely [batch_size, num_tokens, 1, embedding_dim] if batch_feat_with_tokens.dim() == 3: reshaped_batch_feat = batch_feat_with_tokens.squeeze(1) else: reshaped_batch_feat = batch_feat_with_tokens enhanced_batch_feat_with_tokens = self.cross_attention_transformer( reshaped_batch_feat.unsqueeze(0), # Add batch dimension [1, num_tokens, embedding_dim] img_embed ) # Remove batch dimension enhanced_batch_feat_with_tokens = enhanced_batch_feat_with_tokens.squeeze(0) # Reshape back to original shape if needed if batch_feat_with_tokens.dim() == 2: enhanced_batch_feat_with_tokens = enhanced_batch_feat_with_tokens.unsqueeze(1) # Skip connection batch_feat_with_tokens = original_batch_feat_with_tokens + enhanced_batch_feat_with_tokens # print(f"Batch feat with tokens shape: {batch_feat_with_tokens.shape}") # Process all prompts for this image through SAM prompt encoder sparse_embeddings, dense_embeddings = self.visual_model.sam_prompt_encoder( points=None, boxes=None, masks=None, text_embeds=batch_feat_with_tokens, ) sparse_embeddings = sparse_embeddings.to(batch_feat_with_tokens.dtype) high_res_features = [ feat_level[img_idx].unsqueeze(0) for feat_level in _features["high_res_feats"] ] # Process all prompts for this image through SAM mask decoder low_res_masks, iou_pred, _, _ = self.visual_model.sam_mask_decoder( image_embeddings=_features["image_embed"][img_idx].unsqueeze(0), image_pe=self.visual_model.sam_prompt_encoder.get_dense_pe(), sparse_prompt_embeddings=sparse_embeddings, dense_prompt_embeddings=dense_embeddings, multimask_output=False, repeat_image=True, high_res_features=high_res_features, ) # Get predictions for this image pred_masks = low_res_masks.squeeze(1) outputs = {"pred_masks": pred_masks.unsqueeze(0), "pred_logits": iou_pred.unsqueeze(0)} ################################# Inference Postprocessing ################################## # Postprocess masks if not self.training: unique_categories = batched_inputs[img_idx]["unique_categories"] # Assign class labels before filtering num_total_masks = len(pred_masks) # Each unique category gets num_tokens number of predictions class_indices = torch.div(torch.arange(num_total_masks, device=self.device), self.num_tokens, rounding_mode='floor') # Map to actual category IDs from unique_categories class_labels = torch.tensor([unique_categories[i] for i in class_indices], dtype=torch.int64, device=self.device) # FIRST STAGE FILTERING: Filter out low IoU predictions before second stage pred_logits = outputs["pred_logits"].squeeze(0) iou_scores = pred_logits.squeeze(1) if pred_logits.dim() > 1 else pred_logits # Only apply two-stage inference if enabled if self.two_stage_inference: # Apply IoU threshold to filter masks keep_indices = iou_scores >= self.iou_threshold if keep_indices.sum() > 0: # Filter masks based on IoU scores filtered_masks = low_res_masks[keep_indices] filtered_class_labels = class_labels[keep_indices] # SECOND STAGE: Use filtered masks as visual prompts for SAM sparse_embeddings, dense_embeddings = self.visual_model.sam_prompt_encoder( points=None, boxes=None, masks=filtered_masks, text_embeds=None, ) refined_masks, refined_iou_pred, refined_tokens_out, _ = self.visual_model.sam_mask_decoder( image_embeddings=_features["image_embed"][img_idx].unsqueeze(0), image_pe=self.visual_model.sam_prompt_encoder.get_dense_pe(), sparse_prompt_embeddings=sparse_embeddings, dense_prompt_embeddings=dense_embeddings, multimask_output=False, repeat_image=True, high_res_features=high_res_features, ) # Update low_res_masks and outputs with refined predictions low_res_masks = refined_masks pred_logits = refined_iou_pred class_labels = filtered_class_labels # Proceed with postprocessing using the refined masks pred_masks = self.postprocess_masks(low_res_masks, orig_hw=original_size_list[img_idx]) processed_results.append({}) if self.refer_on: # Get referring expression masks refer_masks, refer_scores = self.refer_inference(pred_masks, pred_logits, class_labels) processed_results[-1]["grounding_mask"] = refer_masks processed_results[-1]["grounding_scores"] = refer_scores if self.instance_on: # Process all predictions and perform NMS prompt_results = self.instance_inference(pred_masks, pred_logits, class_labels) # Add instance segmentation results processed_results[-1]["instances"] = prompt_results if self.panoptic_on: # Generate panoptic segmentation directly from predictions # No need to rely on instance results panoptic_r = self.panoptic_inference( pred_logits, # [num_queries, 1] pred_masks, # [num_queries, 1, H, W] class_labels # [num_queries] ) processed_results[-1]["panoptic_seg"] = panoptic_r if self.semantic_on: # Prepare inputs for semantic inference # Create one-hot class scores num_classes = len(self.metadata.stuff_classes) mask_cls = torch.zeros((pred_masks.shape[0], num_classes + 1), device=self.device) # +1 for background # Fill in class scores based on class labels and prediction scores for idx, (cls_id, score) in enumerate(zip(class_labels, pred_logits.squeeze(1))): mask_cls[idx, cls_id] = score # Generate semantic segmentation sem_seg = self.semantic_inference(mask_cls, pred_masks, keep_sem_bgd=False) processed_results[-1]["sem_seg"] = sem_seg return processed_results ################################# Calculate Losses ####################################### # Calculate loss for this image if in training mode if self.training: gt_instances = batched_inputs[img_idx]["instances"] if not isinstance(gt_instances, list): gt_instances = [gt_instances] # For per-prompt matching, we need to split the predictions by prompt num_prompts = len(gt_instances) # Each prompt gets self.num_tokens predictions pred_splits = [self.num_tokens] * num_prompts pred_masks_list = torch.split(pred_masks, pred_splits) pred_logits_list = torch.split(iou_pred, pred_splits) # Process each prompt separately for prompt_idx in range(num_prompts): # Create outputs for this prompt prompt_outputs = { "pred_masks": pred_masks_list[prompt_idx].unsqueeze(0), "pred_logits": pred_logits_list[prompt_idx].unsqueeze(0) } # Prepare targets for this prompt prompt_targets = self.prepare_targets([gt_instances[prompt_idx]]) if return_intermediate and prompt_idx == 0: return prompt_outputs, prompt_targets # Calculate losses for this prompt prompt_losses = self.criterion(prompt_outputs, prompt_targets) # Store weighted losses for k, v in prompt_losses.items(): if k in self.criterion.weight_dict: all_losses[k].append(v * self.criterion.weight_dict[k]) # Average losses across batch if self.training: final_losses = {k: torch.stack(v).mean() for k, v in all_losses.items()} return final_losses def prepare_targets(self, targets): new_targets = [] for targets_per_image in targets: gt_masks = targets_per_image.gt_masks.to(dtype=self.dtype, device=self.device) # unlike traditional instance segmentation model that predicts for every instance, # we only want instances that correspond to the prompt queries (conditional predictions), # so we set the labels to 0 for all instances (label doesn't matter for conditional predictions) labels = torch.zeros_like(targets_per_image.gt_classes).to(device=self.device) target_dict = { "labels": labels, "masks": gt_masks, } new_targets.append(target_dict) return new_targets def instance_inference(self, pred_masks, iou_scores, class_labels): """ Postprocess predicted masks and IoU scores to generate instance segmentation results. Args: pred_masks (Tensor): Predicted masks of shape [num_queries, H, W]. iou_scores (Tensor): IoU scores of shape [num_queries, 1]. class_labels (Tensor): Class labels of shape [num_queries]. Returns: Instances: An `Instances` object containing the final masks, boxes, scores, and class IDs. """ test_topk_per_image = self.test_topk_per_image nms_threshold = self.nms_threshold iou_threshold = self.iou_threshold # Filtering IoU threshold top_k = self.top_k_on nms = self.nms_on image_size = pred_masks.shape[-2:] iou_scores = iou_scores.squeeze(1) # Shape: [num_queries] pred_masks = pred_masks.squeeze(1) # Shape: [num_queries, H, W] 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(iou_scores).bool() for i, lab in enumerate(class_labels): keep[i] = lab in thing_dataset_id_to_contiguous_id.values() pred_masks = pred_masks[keep] iou_scores = iou_scores[keep] class_labels = class_labels[keep] # Step 1: Select top-k masks based on IoU scores if top_k: top_k = min(test_topk_per_image, pred_masks.shape[0]) # Ensure top_k does not exceed the number of masks top_k_indices = torch.argsort(iou_scores, descending=True)[:top_k] pred_masks = pred_masks[top_k_indices] iou_scores = iou_scores[top_k_indices] class_labels = class_labels[top_k_indices] # Step 2: Filter masks based on IoU threshold keep_indices = iou_scores >= iou_threshold pred_masks = pred_masks[keep_indices] iou_scores = iou_scores[keep_indices] class_labels = class_labels[keep_indices] if pred_masks.shape[0] == 0: # No valid masks remain after filtering print("No valid masks remain after filtering. Returning an empty Instances object.") # Return an empty Instances object result = Instances(image_size) result.pred_masks = torch.empty((0, image_size[0], image_size[1]), device=self.device) result.pred_boxes = Boxes(torch.empty((0, 4), device=self.device)) result.scores = torch.empty((0,), device=self.device) result.pred_classes = torch.empty((0,), dtype=torch.int64, device=self.device) return result # Step 3: Compute bounding boxes from masks bit_masks = BitMasks(pred_masks > 0) # Binarize masks pred_boxes = bit_masks.get_bounding_boxes().to(device=self.device) # Shape: [num_instances, 4] # Step 4: Non-Maximum Suppression (NMS) if nms: nms_keep = torchvision.ops.nms(pred_boxes.tensor, iou_scores, nms_threshold) pred_masks = pred_masks[nms_keep] pred_boxes = pred_boxes[nms_keep] iou_scores = iou_scores[nms_keep] class_labels = class_labels[nms_keep] # Step 5: Create Instances result = Instances(image_size) result.pred_masks = (pred_masks > 0).float() result.pred_boxes = pred_boxes result.scores = iou_scores result.pred_classes = class_labels return result def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor: """ Perform PostProcessing on output masks. """ masks = masks.float() masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False) return masks def semantic_inference(self, mask_cls, mask_pred, keep_sem_bgd=False): """ Compute semantic segmentation predictions from class scores and predicted masks. Args: mask_cls (Tensor): Class logits of shape [num_queries, num_classes]. mask_pred (Tensor): Binary mask logits of shape [num_queries, H, W]. keep_sem_bgd (bool): Whether to keep background class or not. Returns: Tensor: Semantic segmentation of shape [num_classes, H, W]. """ if keep_sem_bgd: mask_cls = F.softmax(mask_cls, dim=-1) else: mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1] # Remove background class # mask_pred = mask_pred.sigmoid() mask_pred = mask_pred.sigmoid() mask_pred = mask_pred.squeeze(1) semseg = torch.einsum("qc,qhw->chw", mask_cls, mask_pred) return semseg def mask_nms(self, masks, scores, iou_threshold=0.5): """ Apply Non-Maximum Suppression to masks based on their IoU and scores. Args: masks (Tensor): Binary masks of shape [N, H, W] scores (Tensor): Confidence scores of shape [N] iou_threshold (float): IoU threshold for suppression Returns: Tensor: Boolean tensor of shape [N] indicating which masks to keep """ n = masks.shape[0] if n == 0: return torch.zeros(0, dtype=torch.bool, device=masks.device) if n == 1: return torch.ones(1, dtype=torch.bool, device=masks.device) # Ensure masks are binary binary_masks = masks >= 0.5 # Calculate areas of each mask areas = binary_masks.sum(dim=(1, 2)) # Sort by score order = torch.argsort(scores, descending=True) keep = torch.ones(n, dtype=torch.bool, device=masks.device) for i in range(n): # Skip if this mask is already suppressed if not keep[order[i]]: continue # Get the current mask mask_i = binary_masks[order[i]] area_i = areas[order[i]] # Check against all lower-scored masks for j in range(i + 1, n): if not keep[order[j]]: continue # Calculate IoU mask_j = binary_masks[order[j]] area_j = areas[order[j]] intersection = (mask_i & mask_j).sum() union = area_i + area_j - intersection iou = intersection / union if union > 0 else 0 # Suppress mask_j if IoU is above threshold if iou > iou_threshold: keep[order[j]] = False return keep def panoptic_inference(self, mask_cls, mask_pred, class_labels): """ Compute panoptic segmentation predictions from class scores and predicted masks. Args: mask_cls (Tensor): Class confidence scores of shape [num_queries, 1] mask_pred (Tensor): Binary masks of shape [num_queries, H, W] class_labels (Tensor): Class labels of shape [num_queries] Returns: Tuple: (panoptic_seg, segments_info) - panoptic_seg (Tensor): Panoptic segmentation of shape [H, W] - segments_info (List[Dict]): List of dictionaries containing information about each segment """ scores = mask_cls.squeeze(1) # [num_queries] mask_pred = mask_pred.squeeze(1) mask_pred = mask_pred.sigmoid() # [num_queries, H, W] # Filter based on score threshold keep = scores > self.iou_threshold cur_scores = scores[keep] cur_classes = class_labels[keep] cur_masks = mask_pred[keep] # Get image dimensions h, w = cur_masks.shape[-2:] # Initialize panoptic segmentation tensor panoptic_seg = torch.zeros((h, w), dtype=torch.int32, device=self.device) segments_info = [] if cur_masks.shape[0] == 0: # We didn't detect any mask return panoptic_seg, segments_info # Apply NMS per class to remove duplicate predictions class_ids = torch.unique(cur_classes) nms_keep = torch.zeros_like(cur_scores, dtype=torch.bool) for cls_id in class_ids: # Find all masks for this class cls_mask = cur_classes == cls_id if cls_mask.sum() <= 1: # If only one mask for this class, keep it nms_keep[cls_mask] = True continue # Apply NMS to masks of this class cls_keep = self.mask_nms( cur_masks[cls_mask], cur_scores[cls_mask], iou_threshold=self.nms_threshold # NMS IoU threshold ) # Update the overall keep mask nms_keep[torch.where(cls_mask)[0][cls_keep]] = True # Apply NMS filtering cur_scores = cur_scores[nms_keep] cur_classes = cur_classes[nms_keep] cur_masks = cur_masks[nms_keep] # Calculate probabilities for each mask cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks # Take argmax to determine which mask has highest probability at each pixel cur_mask_ids = cur_prob_masks.argmax(0) # Track stuff (non-thing) regions to merge them stuff_memory_list = {} # Get information about which classes are "things" vs. "stuff" thing_dataset_id_to_contiguous_id = {} if hasattr(self.metadata, 'thing_dataset_id_to_contiguous_id'): thing_dataset_id_to_contiguous_id = self.metadata.thing_dataset_id_to_contiguous_id # Process each mask current_segment_id = 0 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() # Get mask area statistics 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) # Skip masks with small valid areas or overlap issues # Use a more relaxed threshold since we've already handled duplicates with NMS if mask_area > 0 and original_area > 0 and mask.sum().item() > 0: if mask_area / original_area < 0.5: # Relaxed from 0.8 to 0.5 continue # Merge stuff regions with same class 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 # Update panoptic segmentation current_segment_id += 1 panoptic_seg[mask] = current_segment_id # Add segment info seg_info = { "id": current_segment_id, "isthing": bool(isthing), "category_id": int(pred_class), } segments_info.append(seg_info) return panoptic_seg, segments_info def refer_inference(self, pred_masks, pred_logits, class_labels): """ For each class, identify the mask prediction that has the highest confidence score. Args: pred_masks (Tensor): Predicted masks of shape [num_queries, H, W] pred_logits (Tensor): Confidence scores of shape [num_queries, 1] class_labels (Tensor): Class labels of shape [num_queries] Returns: Tensor: Mask predictions of shape [num_classes, H, W] """ # Get unique class labels unique_classes = torch.unique(class_labels) num_classes = len(unique_classes) h, w = pred_masks.shape[-2:] # Initialize output tensor class_masks = torch.zeros((num_classes, h, w), device=self.device) class_scores = torch.zeros((num_classes), device=self.device) # For each class, find the mask with highest confidence for i, cls in enumerate(unique_classes): # Get indices for this class cls_indices = (class_labels == cls) if cls_indices.sum() > 0: # Get masks and scores for this class cls_masks = pred_masks[cls_indices] cls_scores = pred_logits[cls_indices].squeeze(-1) # Find mask with highest score best_idx = torch.argmax(cls_scores) best_mask = cls_masks[best_idx] best_score = cls_scores[best_idx] # Store in output tensor class_masks[i] = best_mask class_scores[i] = best_score return class_masks, class_scores # Add the CrossAttentionTransformer class after the OpenWorldSAM2 class definition class CrossAttentionTransformer(nn.Module): """ A stack of Transformer blocks for cross-attention between VLM features and image embeddings. """ def __init__( self, embedding_dim: int, num_heads: int, mlp_dim: int, num_layers: int = 3, # Added parameter for number of layers dropout: float = 0.1, ): super().__init__() self.embedding_dim = embedding_dim self.num_heads = num_heads self.num_layers = num_layers # Create a stack of transformer layers self.layers = nn.ModuleList([ CrossAttentionLayer( embedding_dim=embedding_dim, num_heads=num_heads, mlp_dim=mlp_dim, dropout=dropout ) for _ in range(num_layers) ]) # Add projection layers to handle dimension mismatches self.input_projection = None self.image_projection = None def forward( self, vlm_features: torch.Tensor, # [batch_size, num_tokens, embedding_dim] image_embeddings: torch.Tensor, # [batch_size, H*W, embedding_dim] ) -> torch.Tensor: """ Forward pass through multiple layers of cross-attention. Args: vlm_features: Tensor of shape [batch_size, num_tokens, embedding_dim] image_embeddings: Tensor of shape [batch_size, H*W, embedding_dim] Returns: Tensor of shape [batch_size, num_tokens, embedding_dim] """ # Ensure inputs are 3D tensors with batch dimension assert vlm_features.dim() == 3, f"vlm_features should be 3D, got shape {vlm_features.shape}" assert image_embeddings.dim() == 3, f"image_embeddings should be 3D, got shape {image_embeddings.shape}" # Check if we need to create projection layers for dimension mismatch input_dim = vlm_features.size(-1) image_dim = image_embeddings.size(-1) # Create projection layers if needed and if they don't exist yet if input_dim != self.embedding_dim and self.input_projection is None: print(f"Creating input projection layer from {input_dim} to {self.embedding_dim}") self.input_projection = nn.Linear(input_dim, self.embedding_dim).to(vlm_features.device) if image_dim != self.embedding_dim and self.image_projection is None: print(f"Creating image projection layer from {image_dim} to {self.embedding_dim}") self.image_projection = nn.Linear(image_dim, self.embedding_dim).to(image_embeddings.device) # Apply projections if needed if self.input_projection is not None: vlm_features = self.input_projection(vlm_features) if self.image_projection is not None: image_embeddings = self.image_projection(image_embeddings) # Pass through all layers x = vlm_features for layer in self.layers: x = layer(x, image_embeddings) # Project back to original dimension if needed if self.input_projection is not None: # Create a projection back to the original dimension if not hasattr(self, 'output_projection') or self.output_projection is None: print(f"Creating output projection layer from {self.embedding_dim} to {input_dim}") self.output_projection = nn.Linear(self.embedding_dim, input_dim).to(x.device) x = self.output_projection(x) return x class CrossAttentionLayer(nn.Module): """ A single Transformer layer for cross-attention between VLM features and image embeddings. """ def __init__( self, embedding_dim: int, num_heads: int, mlp_dim: int, dropout: float = 0.1, ): super().__init__() self.embedding_dim = embedding_dim self.num_heads = num_heads # Self-attention for VLM features self.self_attn_norm = nn.LayerNorm(embedding_dim) self.self_attn = nn.MultiheadAttention( embedding_dim, num_heads, dropout=dropout, batch_first=True ) self.self_attn_dropout = nn.Dropout(dropout) # Cross-attention from VLM features to image embeddings self.cross_attn_norm = nn.LayerNorm(embedding_dim) self.cross_attn = nn.MultiheadAttention( embedding_dim, num_heads, dropout=dropout, batch_first=True ) self.cross_attn_dropout = nn.Dropout(dropout) # MLP block self.mlp_norm = nn.LayerNorm(embedding_dim) self.mlp = nn.Sequential( nn.Linear(embedding_dim, mlp_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(mlp_dim, embedding_dim), nn.Dropout(dropout) ) def forward( self, vlm_features: torch.Tensor, # [batch_size, num_tokens, embedding_dim] image_embeddings: torch.Tensor, # [batch_size, H*W, embedding_dim] ) -> torch.Tensor: """ Forward pass for a single cross-attention layer. Args: vlm_features: Tensor of shape [batch_size, num_tokens, embedding_dim] image_embeddings: Tensor of shape [batch_size, H*W, embedding_dim] Returns: Tensor of shape [batch_size, num_tokens, embedding_dim] """ # Self-attention residual = vlm_features x = self.self_attn_norm(vlm_features) x, _ = self.self_attn(x, x, x) x = self.self_attn_dropout(x) x = residual + x # Cross-attention residual = x x = self.cross_attn_norm(x) x, _ = self.cross_attn( query=x, key=image_embeddings, value=image_embeddings ) x = self.cross_attn_dropout(x) x = residual + x # MLP residual = x x = self.mlp_norm(x) x = self.mlp(x) x = residual + x return x