Add CLIP-based semantic re-ranking for instance segmentation

configuration_openworld_sam.py

@@ -53,6 +53,21 @@ class OpenWorldSAMConfig(PretrainedConfig):
         iou_thresh=0.7,
         top_k_on=False,
         detections_per_image=100,
+        # Mask decoder batching: the decoder's repeat_image path duplicates
+        # the image embeddings/positional encoding once per candidate query,
+        # so running all num_classes * num_tokens candidates through it in
+        # one call can reach double-digit GiB for large vocabularies. This
+        # bounds peak memory by running the decoder in chunks instead.
+        mask_decoder_chunk_size=500,
+        # CLIP-based re-ranking: the SAM mask-quality (IoU) score measures
+        # only whether a mask is well-formed, not whether it matches the
+        # queried class, so class-agnostic mask proposals conditioned on
+        # different prompts routinely tie on score for the same object. When
+        # enabled, surviving candidates (post iou_thresh) are re-scored and
+        # re-labeled using actual CLIP image/text similarity instead.
+        use_clip_rerank=True,
+        clip_model_name_or_path="openai/clip-vit-base-patch32",
+        clip_crop_padding_frac=0.1,
         # Class vocabulary
         stuff_classes=None,
         **kwargs,
@@ -74,4 +89,8 @@ class OpenWorldSAMConfig(PretrainedConfig):
         self.iou_thresh = iou_thresh
         self.top_k_on = top_k_on
         self.detections_per_image = detections_per_image
+        self.mask_decoder_chunk_size = mask_decoder_chunk_size
+        self.use_clip_rerank = use_clip_rerank
+        self.clip_model_name_or_path = clip_model_name_or_path
+        self.clip_crop_padding_frac = clip_crop_padding_frac
         self.stuff_classes = stuff_classes if stuff_classes is not None else ADE20K_150_CLASSES

modeling_openworld_sam.py

@@ -180,6 +180,10 @@ class OpenWorldSAMModel(PreTrainedModel):
         # Tokenizer loaded lazily on first forward
         self._tokenizer = None
 
+        # CLIP model/processor (semantic re-ranking) loaded lazily on first use
+        self._clip_model = None
+        self._clip_processor = None
+
         # Required by transformers>=5's from_pretrained (sets
         # self.all_tied_weights_keys and other bookkeeping consumed by
         # _finalize_model_loading); harmless no-op pre-checkpoint-load init.
@@ -199,6 +203,31 @@ class OpenWorldSAMModel(PreTrainedModel):
             )
         return self._tokenizer
 
+    # ------------------------------------------------------------------
+    # CLIP (semantic re-ranking)
+    # ------------------------------------------------------------------
+
+    @property
+    def clip_model(self):
+        if self._clip_model is None:
+            from transformers import CLIPModel
+
+            self._clip_model = CLIPModel.from_pretrained(
+                self.config.clip_model_name_or_path
+            ).to(self.device)
+            self._clip_model.eval()
+        return self._clip_model
+
+    @property
+    def clip_processor(self):
+        if self._clip_processor is None:
+            from transformers import CLIPProcessor
+
+            self._clip_processor = CLIPProcessor.from_pretrained(
+                self.config.clip_model_name_or_path
+            )
+        return self._clip_processor
+
     def _tokenize_prompts(self, prompts):
         tok = self.tokenizer
         ids = [tok(p, return_tensors="pt").input_ids[0] for p in prompts]
@@ -388,11 +417,11 @@ class OpenWorldSAMModel(PreTrainedModel):
         iou_pred_all = torch.cat(iou_pred_chunks, dim=0)
 
         pred_masks_all = low_res_masks_all.squeeze(1)  # [total_tokens_all, H_low, W_low]
-        pred_logits_all = iou_pred_all.squeeze(1)      # [total_tokens_all]
-        
+        pred_logits_all = iou_pred_all.squeeze(1)       # [total_tokens_all]
+
         pred_masks_per_image = torch.split(pred_masks_all, token_counts)
         pred_logits_per_image = torch.split(pred_logits_all, token_counts)
-                
+
         processed_results = []
         for img_idx in range(batch_size):
             pred_masks = pred_masks_per_image[img_idx]
@@ -410,14 +439,31 @@ class OpenWorldSAMModel(PreTrainedModel):
                 dtype=torch.long, device=self.device,
             )
 
+            square_rgb = None
+            id_to_text = None
+            if self.config.use_clip_rerank:
+                # Reads from config directly rather than the pixel_mean/std
+                # buffers: those are registered persistent=False (excluded
+                # from the checkpoint on purpose), but transformers'
+                # from_pretrained fast-init path zeroes any buffer it
+                # doesn't find a matching checkpoint key for, config values
+                # aren't affected by that.
+                mean = torch.tensor(self.config.pixel_mean, device=images.device).view(-1, 1, 1)
+                std = torch.tensor(self.config.pixel_std, device=images.device).view(-1, 1, 1)
+                denorm = images[img_idx].float() * std + mean
+                square_rgb = denorm.clamp(0, 255).round().byte().permute(1, 2, 0).cpu().numpy()
+                id_to_text = dict(zip(
+                    unique_categories, batched_inputs[img_idx]["prompt"]
+                ))
+
             # Filter on low-res masks first; only the survivors get upsampled
             # to the original image size (upsampling the full query set before
             # filtering allocates one [num_queries, H, W] float32 tensor that
             # can reach double-digit GiB for large vocabularies).
             instances = self._instance_inference(
-                pred_masks, pred_logits, class_labels, original_size_list[img_idx]
+                pred_masks, pred_logits, class_labels, original_size_list[img_idx],
+                square_rgb=square_rgb, id_to_text=id_to_text,
             )
-
             processed_results.append({"instances": instances})
 
         return processed_results
@@ -427,7 +473,93 @@ class OpenWorldSAMModel(PreTrainedModel):
             masks.float().unsqueeze(0), orig_hw, mode="bilinear", align_corners=False
         ).squeeze(0)
 
-    def _instance_inference(self, pred_masks, iou_scores, class_labels, orig_hw):
+    def _clip_rerank(self, pred_masks, class_labels, square_rgb, id_to_text):
+        """Re-scores/re-labels candidates using real CLIP image/text similarity.
+
+        The SAM mask-quality (IoU) score measures whether a mask is well
+        formed, not whether it matches the queried class: each prompt is run
+        through cross-attention independently, so several unrelated prompts
+        routinely produce near-identical masks over the same salient object,
+        each with a similarly high IoU score. This crops each surviving
+        candidate's mask region out of the original image and asks a real
+        (frozen, off-the-shelf) CLIP model which of the candidate class names
+        actually matches that crop, replacing both the class label and the
+        score with CLIP's pick/confidence.
+
+        Args:
+            pred_masks: bool/float tensor [N, h, w], low-res mask logits/mask
+            class_labels: long tensor [N], the prompt each mask was proposed under
+            square_rgb: uint8 numpy array [H, W, 3], the (square) SAM input image
+            id_to_text: dict mapping category id -> prompt string
+
+        Returns:
+            (class_labels, scores) tensors, both length N
+        """
+        n = pred_masks.shape[0]
+        if n == 0:
+            return class_labels, torch.empty(0, device=self.device)
+
+        from PIL import Image
+
+        H, W = square_rgb.shape[:2]
+        masks_up = F.interpolate(
+            (pred_masks > 0).float().unsqueeze(1), (H, W), mode="nearest"
+        ).squeeze(1).bool()
+
+        padding_frac = self.config.clip_crop_padding_frac
+        crops = []
+        for i in range(n):
+            ys, xs = torch.where(masks_up[i])
+            if ys.numel() == 0:
+                crops.append(None)
+                continue
+            y1, y2 = int(ys.min()), int(ys.max())
+            x1, x2 = int(xs.min()), int(xs.max())
+            ph = int((y2 - y1 + 1) * padding_frac)
+            pw = int((x2 - x1 + 1) * padding_frac)
+            y1, y2 = max(0, y1 - ph), min(H, y2 + 1 + ph)
+            x1, x2 = max(0, x1 - pw), min(W, x2 + 1 + pw)
+            crops.append(Image.fromarray(square_rgb[y1:y2, x1:x2]))
+
+        unique_ids = sorted(set(id_to_text.keys()))
+        texts = [id_to_text[c] for c in unique_ids]
+
+        valid_idx = [i for i, c in enumerate(crops) if c is not None]
+        new_class_labels = class_labels.clone()
+        new_scores = torch.zeros(n, device=self.device)
+        if not valid_idx:
+            return new_class_labels, new_scores
+
+        with torch.no_grad():
+            text_inputs = self.clip_processor(
+                text=texts, return_tensors="pt", padding=True
+            ).to(self.device)
+            # get_text_features/get_image_features return a
+            # BaseModelOutputWithPooling whose `.pooler_output` has been
+            # overwritten with the projected embedding, not a raw tensor.
+            text_feats = self.clip_model.get_text_features(**text_inputs).pooler_output
+            text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
+
+            img_inputs = self.clip_processor(
+                images=[crops[i] for i in valid_idx], return_tensors="pt"
+            ).to(self.device)
+            img_feats = self.clip_model.get_image_features(**img_inputs).pooler_output
+            img_feats = img_feats / img_feats.norm(dim=-1, keepdim=True)
+
+            logit_scale = self.clip_model.logit_scale.exp()
+            probs = ((img_feats @ text_feats.T) * logit_scale).softmax(dim=-1)
+
+        for row, i in enumerate(valid_idx):
+            best = int(probs[row].argmax())
+            new_class_labels[i] = unique_ids[best]
+            new_scores[i] = probs[row, best]
+
+        return new_class_labels, new_scores
+
+    def _instance_inference(
+        self, pred_masks, iou_scores, class_labels, orig_hw,
+        square_rgb=None, id_to_text=None,
+    ):
         """Returns dict with keys: masks (bool), scores (float), class_ids (long)."""
         pred_masks = pred_masks.squeeze(1) if pred_masks.ndim == 4 else pred_masks
 
@@ -449,6 +581,12 @@ class OpenWorldSAMModel(PreTrainedModel):
                 "class_ids": empty.long(),
             }
 
+        # Re-score/re-label using real CLIP image/text similarity instead of
+        # the SAM mask-quality score, which doesn't reflect class match.
+        if self.config.use_clip_rerank and square_rgb is not None and id_to_text is not None:
+            class_labels, iou_scores = self._clip_rerank(
+                pred_masks, class_labels, square_rgb, id_to_text
+            )
 
         # NMS on low-res masks — box IoU is scale-invariant, so this doesn't
         # need the full-res masks either.