| """Patch aggregation and mask-aware region encoding for Falcon. |
| |
| The detector contract used here is deliberately small: boxes are ``xyxy`` in |
| original-image coordinates and masks are either logits or probabilities. The |
| module never guesses which representation a mask uses and never applies a |
| sigmoid; it makes a deterministic binary mask with the declared threshold. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| from torch import Tensor, nn |
| from torch.nn import functional as F |
| from torchvision.ops import roi_align |
|
|
|
|
| class PatchTokenAggregator(nn.Module): |
| """Concatenate each non-overlapping 2x2 patch block, then project it. |
| |
| DINOv2 emits patch tokens in row-major order. For a block, concatenation |
| order is top-left, top-right, bottom-left, bottom-right. |
| """ |
|
|
| def __init__(self, input_dim: int, output_dim: int) -> None: |
| super().__init__() |
| self.input_dim = input_dim |
| self.output_dim = output_dim |
| self.projection = nn.Linear(4 * input_dim, output_dim) |
|
|
| def forward(self, patch_tokens: Tensor, grid_size: tuple[int, int]) -> Tensor: |
| if patch_tokens.ndim != 3: |
| raise ValueError("patch_tokens must have shape [batch, patches, channels]") |
| batch, count, channels = patch_tokens.shape |
| height, width = grid_size |
| if channels != self.input_dim: |
| raise ValueError(f"expected {self.input_dim} channels, received {channels}") |
| if count != height * width: |
| raise ValueError(f"{count} tokens cannot form the declared {height}x{width} grid") |
| if height % 2 or width % 2: |
| raise ValueError("the DINOv2 patch grid must be even for 2x2 aggregation") |
|
|
| grid = patch_tokens.reshape(batch, height // 2, 2, width // 2, 2, channels) |
| blocks = grid.permute(0, 1, 3, 2, 4, 5).reshape( |
| batch, height // 2, width // 2, 4 * channels |
| ) |
| return self.projection(blocks).permute(0, 3, 1, 2).contiguous() |
|
|
|
|
| class MaskAwareRegionEncoder(nn.Module): |
| """Fuse ROI-aligned and binary-mask-pooled features by concatenation. |
| |
| Invalid/padded proposals are returned as exact zero vectors. ``image_sizes`` |
| contains ``(height, width)`` for each original detector image. |
| """ |
|
|
| def __init__( |
| self, |
| feature_dim: int, |
| output_dim: int, |
| *, |
| roi_size: int = 4, |
| sampling_ratio: int = 2, |
| probability_threshold: float = 0.5, |
| ) -> None: |
| super().__init__() |
| if not 0.0 <= probability_threshold <= 1.0: |
| raise ValueError("probability_threshold must lie in [0, 1]") |
| self.feature_dim = feature_dim |
| self.output_dim = output_dim |
| self.roi_size = roi_size |
| self.sampling_ratio = sampling_ratio |
| self.probability_threshold = probability_threshold |
| self.projection = nn.Linear(2 * feature_dim, output_dim) |
|
|
| @staticmethod |
| def _validate_inputs( |
| feature_map: Tensor, |
| boxes: Tensor, |
| masks: Tensor, |
| image_sizes: Tensor, |
| valid: Tensor, |
| ) -> None: |
| if feature_map.ndim != 4: |
| raise ValueError("feature_map must have shape [batch, channels, height, width]") |
| if boxes.ndim != 3 or boxes.shape[-1] != 4: |
| raise ValueError("boxes must have shape [batch, proposals, 4]") |
| if masks.ndim != 4 or masks.shape[:2] != boxes.shape[:2]: |
| raise ValueError("masks must have shape [batch, proposals, height, width]") |
| if image_sizes.shape != (feature_map.shape[0], 2): |
| raise ValueError("image_sizes must have shape [batch, 2] in (height, width) order") |
| if valid.shape != boxes.shape[:2]: |
| raise ValueError("valid must have shape [batch, proposals]") |
| if boxes.shape[0] != feature_map.shape[0]: |
| raise ValueError("feature_map and detector outputs must have the same batch size") |
|
|
| def pool_components( |
| self, |
| feature_map: Tensor, |
| boxes: Tensor, |
| masks: Tensor, |
| image_sizes: Tensor, |
| valid: Tensor | None = None, |
| *, |
| masks_are_logits: bool = False, |
| mask_sizes: Tensor | None = None, |
| ) -> tuple[Tensor, Tensor, Tensor]: |
| """Return valid proposal indices, ROI vectors, and mask vectors.""" |
|
|
| if valid is None: |
| valid = torch.ones(boxes.shape[:2], dtype=torch.bool, device=boxes.device) |
| else: |
| valid = valid.to(dtype=torch.bool) |
| self._validate_inputs(feature_map, boxes, masks, image_sizes, valid) |
|
|
| indices = valid.nonzero(as_tuple=False) |
| if indices.numel() == 0: |
| empty = feature_map.new_empty((0, self.feature_dim)) |
| return indices, empty, empty |
|
|
| device = feature_map.device |
| dtype = feature_map.dtype |
| indices = indices.to(device=device) |
| boxes = boxes.to(device=device, dtype=dtype) |
| masks = masks.to(device=device) |
| image_sizes = image_sizes.to(device=device, dtype=dtype) |
|
|
| batch_ids, proposal_ids = indices.unbind(dim=1) |
| selected_boxes = boxes[batch_ids, proposal_ids].clone() |
| source_h = image_sizes[batch_ids, 0].clamp_min(1) |
| source_w = image_sizes[batch_ids, 1].clamp_min(1) |
| feature_h, feature_w = feature_map.shape[-2:] |
| selected_boxes[:, 0::2] *= feature_w / source_w[:, None] |
| selected_boxes[:, 1::2] *= feature_h / source_h[:, None] |
| rois = torch.cat((batch_ids.to(dtype=dtype).unsqueeze(1), selected_boxes), dim=1) |
|
|
| roi_features = roi_align( |
| feature_map, |
| rois, |
| output_size=(self.roi_size, self.roi_size), |
| spatial_scale=1.0, |
| sampling_ratio=self.sampling_ratio, |
| aligned=True, |
| ) |
| roi_vectors = roi_features.mean(dim=(-2, -1)) |
|
|
| threshold = 0.0 if masks_are_logits else self.probability_threshold |
| if mask_sizes is None: |
| mask_sizes = torch.tensor([list(masks.shape[-2:])] * feature_map.shape[0]) |
| if mask_sizes.shape != (feature_map.shape[0], 2): |
| raise ValueError("mask_sizes must have shape [batch, 2]") |
| binary_masks = feature_map.new_zeros((len(indices), 1, feature_h, feature_w)) |
| for batch_index in range(feature_map.shape[0]): |
| selected = batch_ids == batch_index |
| if not bool(selected.any()): |
| continue |
| height, width = (int(value) for value in mask_sizes[batch_index]) |
| if not (0 < height <= masks.shape[-2] and 0 < width <= masks.shape[-1]): |
| raise ValueError("mask_sizes exceed the padded canvas or are nonpositive") |
| selected_masks = masks[batch_index, proposal_ids[selected], :height, :width].unsqueeze( |
| 1 |
| ) |
| binary = ( |
| selected_masks |
| if selected_masks.dtype == torch.bool |
| else selected_masks >= threshold |
| ) |
| binary_masks[selected] = F.interpolate( |
| binary.to(dtype=dtype), size=(feature_h, feature_w), mode="nearest" |
| ).to(dtype=dtype) |
| selected_features = feature_map.index_select(0, batch_ids) |
| denominator = binary_masks.sum(dim=(-2, -1)).clamp_min(1.0) |
| mask_vectors = (selected_features * binary_masks).sum(dim=(-2, -1)) / denominator |
| return indices, roi_vectors, mask_vectors |
|
|
| def forward( |
| self, |
| feature_map: Tensor, |
| boxes: Tensor, |
| masks: Tensor, |
| image_sizes: Tensor | Sequence[tuple[int, int]], |
| valid: Tensor | None = None, |
| *, |
| masks_are_logits: bool = False, |
| mask_sizes: Tensor | None = None, |
| ) -> Tensor: |
| if not isinstance(image_sizes, Tensor): |
| image_sizes = torch.as_tensor(image_sizes, device=feature_map.device) |
| if valid is None: |
| valid = torch.ones(boxes.shape[:2], dtype=torch.bool, device=boxes.device) |
| indices, roi_vectors, mask_vectors = self.pool_components( |
| feature_map, |
| boxes, |
| masks, |
| image_sizes, |
| valid, |
| masks_are_logits=masks_are_logits, |
| mask_sizes=mask_sizes, |
| ) |
|
|
| batch, proposals = boxes.shape[:2] |
| flat_output = feature_map.new_zeros((batch * proposals, self.output_dim)) |
| if indices.numel() == 0: |
| return flat_output.reshape(batch, proposals, self.output_dim) |
|
|
| fused = self.projection(torch.cat((roi_vectors, mask_vectors), dim=-1)) |
| flat_indices = indices[:, 0] * proposals + indices[:, 1] |
| flat_output = flat_output.index_copy(0, flat_indices, fused) |
| return flat_output.reshape(batch, proposals, self.output_dim) |
|
|
