File size: 8,776 Bytes
c645013 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | """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)
|