From 2efc4d91d7363d65b6f38e7d47c585231bb98158 Mon Sep 17 00:00:00 2001 From: YilinWang121054 <1210543783@qq.com> Date: Fri, 11 Sep 2026 01:27:42 +0800 Subject: [PATCH] =?UTF-8?q?[=E7=8A=80=E7=89=9B=E9=B8=9F-A2]=20Add=20config?= =?UTF-8?q?urable=20adaptive=20STAL?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidated onto upstream 7bfbfd3. The complete tree matches local tested merge 6f55860. Training evidence retains its original source commits. --- tests/test_stal_assignment.py | 212 ++++++++++++++++++++++++++++++++++ ultralytics/cfg/__init__.py | 53 ++++++++- ultralytics/cfg/default.yaml | 10 ++ ultralytics/utils/loss.py | 8 ++ ultralytics/utils/tal.py | 104 +++++++++++++++-- 5 files changed, 372 insertions(+), 15 deletions(-) create mode 100644 tests/test_stal_assignment.py diff --git a/tests/test_stal_assignment.py b/tests/test_stal_assignment.py new file mode 100644 index 000000000..5b7964a03 --- /dev/null +++ b/tests/test_stal_assignment.py @@ -0,0 +1,212 @@ +"""Tests for small-target adaptive task-aligned assignment.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from ultralytics.cfg import get_cfg +from ultralytics.utils.tal import TaskAlignedAssigner + + +def _grid_points(size=3, stride=8): + coordinates = torch.arange(size, dtype=torch.float32) * stride + stride / 2 + yy, xx = torch.meshgrid(coordinates, coordinates, indexing="ij") + return torch.stack((xx, yy), dim=-1).reshape(-1, 2) + + +def test_stal_default_config_preserves_fixed_stride_mode(): + """The new configuration contract must retain the repository's existing behavior by default.""" + cfg = get_cfg() + assert cfg.stal_mode == "fixed" + assert cfg.stal_small_area == 32**2 + assert cfg.stal_medium_area == 96**2 + assert cfg.stal_topk_small >= cfg.stal_min_candidates + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"stal_mode": "unknown"}, "stal_mode"), + ({"stal_small_area": 1024.0, "stal_medium_area": 1024.0}, "0 < small < medium"), + ({"stal_candidate_scale": 0.5}, "between 1.0 and 4.0"), + ({"stal_min_candidates": 4, "stal_topk_small": 3}, "greater than or equal"), + ({"stal_topk_small": 1.5}, "stal_topk_small"), + ({"stal_small_area": float("nan")}, "finite"), + ({"stal_medium_area": float("inf")}, "finite"), + ({"stal_candidate_scale": float("nan")}, "finite"), + ], +) +def test_stal_config_rejects_invalid_values_and_relationships(overrides, match): + """STAL experiment parameters must fail fast instead of silently changing an experiment.""" + with pytest.raises((TypeError, ValueError), match=match): + get_cfg(overrides=overrides) + + +def test_detection_loss_receives_stal_configuration(): + """Training configuration must reach the assigner used by detection loss.""" + import ultralytics.nn.tasks # noqa: F401 + from ultralytics.utils.loss import v8DetectionLoss + + class TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1)) + self.args = get_cfg(overrides={"stal_mode": "adaptive", "stal_topk_small": 15}) + self.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=2, reg_max=16)] + + assigner = v8DetectionLoss(TinyModel()).assigner + assert assigner.stal_mode == "adaptive" + assert assigner.stal_topk[0] == 15 + + +def test_tal_mode_uses_unmodified_gt_box(): + """Disabling STAL must be equivalent to raw TAL candidate geometry.""" + points = _grid_points(size=2) + box = torch.tensor([[[7.0, 7.0, 9.0, 9.0]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + mask = TaskAlignedAssigner(stal_mode="tal").select_candidates_in_gts(points, box, valid) + lt, rb = box.unsqueeze(2).chunk(2, 3) + expected = ((points - lt > 1e-9) & (rb - points > 1e-9)).all(3) + torch.testing.assert_close(mask, expected) + assert not mask.any() + + +def test_fixed_mode_keeps_legacy_stride_expansion(): + """The default fixed mode must keep expanding sub-stride dimensions to the middle stride.""" + points = _grid_points(size=2) + box = torch.tensor([[[7.0, 7.0, 9.0, 9.0]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + mask = TaskAlignedAssigner(stal_mode="fixed", stride=[8, 16, 32]).select_candidates_in_gts(points, box, valid) + assert mask.sum().item() == 4 + + +def test_adaptive_area_boundaries_use_coco_style_training_thresholds(): + """Boundary values are small only below 32^2 and large from 96^2 onward.""" + boxes = torch.tensor([[[0.0, 0.0, 31.0, 33.0], [0.0, 0.0, 32.0, 32.0], [0.0, 0.0, 96.0, 96.0]]]) + assigner = TaskAlignedAssigner( + stal_mode="adaptive", stal_min_candidates=1, stal_topk_small=13, stal_topk_medium=7, stal_topk_large=3 + ) + assert assigner.get_adaptive_topks(boxes).tolist() == [[13, 7, 3]] + + +def test_adaptive_mode_guarantees_minimum_pre_conflict_candidates_for_tiny_gt(): + """A tiny valid GT receives the configured number of candidates before conflict resolution.""" + points = _grid_points(size=3) + box = torch.tensor([[[7.5, 7.5, 8.5, 8.5]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + assigner = TaskAlignedAssigner( + stal_mode="adaptive", stal_min_candidates=5, stal_topk_small=5, stal_topk_medium=5, stal_topk_large=5 + ) + mask = assigner.select_candidates_in_gts(points, box, valid) + assert mask.sum().item() >= 5 + assigner.bs, assigner.n_max_boxes = 1, 1 + positive_mask, _, _ = assigner.get_pos_mask( + torch.zeros(1, points.shape[0], 1), + torch.zeros(1, points.shape[0], 4), + torch.zeros(1, 1, 1), + box, + points, + valid, + ) + assert positive_mask.sum().item() >= 5 + + +def test_adaptive_assigner_handles_empty_gt(): + """Empty batches must return correctly shaped finite targets without positives.""" + points = _grid_points(size=2) + anchors = points.shape[0] + outputs = TaskAlignedAssigner(topk=3, num_classes=2, stal_mode="adaptive")( + torch.full((1, anchors, 2), 0.5), + torch.zeros(1, anchors, 4), + points, + torch.zeros(1, 0, 1), + torch.zeros(1, 0, 4), + torch.zeros(1, 0, 1, dtype=torch.bool), + ) + assert not outputs[3].any() + assert all(torch.isfinite(value).all() for value in outputs if value.is_floating_point()) + + +def test_adaptive_assigner_resolves_overlapping_gt_conflicts_once_per_anchor(): + """Overlapping GT candidates must still resolve to at most one target for each anchor.""" + points = _grid_points(size=4) + anchors = points.shape[0] + boxes = torch.tensor([[[4.0, 4.0, 20.0, 20.0], [8.0, 8.0, 24.0, 24.0]]]) + labels = torch.tensor([[[0.0], [1.0]]]) + valid = torch.ones(1, 2, 1, dtype=torch.bool) + pd_boxes = torch.cat((points - 6, points + 6), dim=-1).unsqueeze(0) + pd_scores = torch.linspace(0.55, 0.95, anchors).view(1, anchors, 1).expand(-1, -1, 2) + assigner = TaskAlignedAssigner( + topk=5, + num_classes=2, + stal_mode="adaptive", + stal_min_candidates=3, + stal_topk_small=5, + stal_topk_medium=5, + stal_topk_large=5, + ) + assigner.bs, assigner.n_max_boxes = 1, 2 + mask_pos, align_metric, overlaps = assigner.get_pos_mask(pd_scores, pd_boxes, labels, boxes, points, valid) + _, fg_mask, resolved = assigner.select_highest_overlaps(mask_pos, overlaps, 2, align_metric) + assert (resolved.sum(1) <= 1).all() + assert fg_mask.sum() > 0 + assert torch.isfinite(align_metric).all() + + +def test_conflict_resolution_cannot_assign_anchor_to_non_candidate_gt(): + """Zero-overlap ties must stay within the GTs that proposed the contested anchor.""" + assigner = TaskAlignedAssigner(num_classes=3, stal_mode="adaptive") + mask_pos = torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]]) + overlaps = torch.zeros_like(mask_pos) + align_metric = torch.zeros_like(mask_pos) + original_counts = mask_pos.sum(-1) + + _, fg_mask, resolved = assigner.select_highest_overlaps( + mask_pos, overlaps, n_max_boxes=3, align_metric=align_metric + ) + + assert resolved[0, 0, 2] == 0 + assert (resolved.sum(-1) <= original_counts).all() + assert (fg_mask <= 1).all() + + +def _assignment_and_gradient(amp_enabled): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + amp_dtype = torch.float16 if device.type == "cuda" else torch.bfloat16 + points = _grid_points(size=4).to(device) + anchors = points.shape[0] + logits = torch.linspace(-1.0, 1.0, anchors, device=device).view(1, anchors, 1).requires_grad_() + pd_boxes = torch.cat((points - 7, points + 7), dim=-1).unsqueeze(0) + gt_boxes = torch.tensor([[[4.0, 4.0, 20.0, 20.0]]], device=device) + gt_labels = torch.zeros(1, 1, 1, device=device) + valid = torch.ones(1, 1, 1, dtype=torch.bool, device=device) + assigner = TaskAlignedAssigner( + topk=5, + num_classes=1, + stal_mode="adaptive", + stal_min_candidates=3, + stal_topk_small=5, + stal_topk_medium=5, + stal_topk_large=5, + ).to(device) + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled): + scores = logits.sigmoid() + outputs = assigner(scores.detach(), pd_boxes, points, gt_labels, gt_boxes, valid) + loss = F.binary_cross_entropy_with_logits(logits, outputs[2].to(logits.dtype)) + loss.backward() + return outputs[3].cpu(), outputs[4].cpu(), float(loss.detach()), logits.grad.detach().cpu() + + +def test_adaptive_assigner_fp32_amp_masks_counts_loss_and_gradients_are_stable(): + """FP32 and autocast must keep assignment decisions aligned and all optimization values finite.""" + fp32_mask, fp32_idx, fp32_loss, fp32_grad = _assignment_and_gradient(False) + amp_mask, amp_idx, amp_loss, amp_grad = _assignment_and_gradient(True) + torch.testing.assert_close(amp_mask, fp32_mask) + torch.testing.assert_close(amp_idx[amp_mask], fp32_idx[fp32_mask]) + assert amp_mask.sum() == fp32_mask.sum() + assert fp32_loss == pytest.approx(amp_loss, rel=0.02, abs=1e-4) + assert torch.isfinite(fp32_grad).all() and torch.isfinite(amp_grad).all() diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index de72fb6fd..2d7a42dbd 100644 --- a/ultralytics/cfg/__init__.py +++ b/ultralytics/cfg/__init__.py @@ -5,6 +5,7 @@ import ast import importlib.util import json +import math import os import shutil import subprocess @@ -263,6 +264,7 @@ "sigma", } ) +STAL_FLOAT_KEYS = frozenset({"stal_candidate_scale", "stal_medium_area", "stal_small_area"}) # fmt: off CFG_FLOAT_KEYS = frozenset( { # integer or float arguments, i.e. x=2 and x=2.0 @@ -290,7 +292,7 @@ "workspace", "batch", } -) | MIXTURE_FLOAT_KEYS +) | MIXTURE_FLOAT_KEYS | STAL_FLOAT_KEYS CFG_FRACTION_KEYS = frozenset( { # fractional floats use [0.0, 1.0], except dataset fraction uses (0.0, 1.0] "dropout", @@ -359,6 +361,7 @@ "slice_size", } ) +STAL_INT_KEYS = frozenset({"stal_min_candidates", "stal_topk_large", "stal_topk_medium", "stal_topk_small"}) CFG_INT_KEYS = frozenset( { # integer-only arguments "epochs", @@ -373,7 +376,7 @@ "nbs", "save_period", } -) | MIXTURE_INT_KEYS +) | MIXTURE_INT_KEYS | STAL_INT_KEYS CFG_INT_MIN = { # minimum valid values for integer arguments used as divisors, sizes or seeds "nbs": 1, "max_det": 1, @@ -495,6 +498,7 @@ "mot_scene_inference_mode", } ) +STAL_STR_KEYS = frozenset({"stal_mode"}) CFG_STR_KEYS = frozenset( { "optimizer", @@ -512,13 +516,14 @@ "foundation_dinov3_weights", "foundation_siglip2_weights", } -) | MIXTURE_STR_KEYS +) | MIXTURE_STR_KEYS | STAL_STR_KEYS FOUNDATION_TEACHERS = frozenset({"none", "dinov3", "siglip2", "multi"}) FOUNDATION_BACKENDS = frozenset({"transformers", "local"}) FOUNDATION_LOSSES = frozenset({"cosine", "l2", "relational", "hybrid"}) FOUNDATION_RELATION_MODES = frozenset({"sampled", "full"}) FOUNDATION_DTYPES = frozenset({"auto", "fp32", "fp16", "bf16"}) FOUNDATION_TARGET_LEVELS = frozenset({"p3", "p4", "p5"}) +STAL_MODES = frozenset({"tal", "fixed", "adaptive"}) # fmt: on LORA_RUNTIME_METADATA_KEYS = frozenset( { @@ -631,6 +636,7 @@ def get_cfg( # Type and Value checks check_cfg(cfg) + validate_stal_config(cfg) validate_foundation_config(cfg) # Return instance @@ -750,6 +756,47 @@ def check_cfg(cfg: dict, hard: bool = True) -> None: cfg[k] = scheme +def validate_stal_config(cfg: dict) -> None: + """Validate small-target adaptive label-assignment values and relationships.""" + mode = cfg.get("stal_mode", DEFAULT_CFG_DICT["stal_mode"]) + small_area = cfg.get("stal_small_area", DEFAULT_CFG_DICT["stal_small_area"]) + medium_area = cfg.get("stal_medium_area", DEFAULT_CFG_DICT["stal_medium_area"]) + candidate_scale = cfg.get("stal_candidate_scale", DEFAULT_CFG_DICT["stal_candidate_scale"]) + min_candidates = cfg.get("stal_min_candidates", DEFAULT_CFG_DICT["stal_min_candidates"]) + topks = { + name: cfg.get(name, DEFAULT_CFG_DICT[name]) + for name in ("stal_topk_small", "stal_topk_medium", "stal_topk_large") + } + + if not isinstance(mode, str): + raise TypeError(f"'stal_mode={mode}' is of invalid type {type(mode).__name__}. 'stal_mode' must be a str.") + if mode not in STAL_MODES: + raise ValueError(f"'stal_mode={mode}' is invalid. Valid values are {sorted(STAL_MODES)}.") + + for name, value in { + "stal_small_area": small_area, + "stal_medium_area": medium_area, + "stal_candidate_scale": candidate_scale, + }.items(): + if isinstance(value, bool) or not isinstance(value, FLOAT_OR_INT): + raise TypeError(f"'{name}={value}' must be an int or float.") + if not math.isfinite(value): + raise ValueError(f"'{name}={value}' must be finite.") + if small_area <= 0 or medium_area <= 0 or small_area >= medium_area: + raise ValueError("'stal_small_area' and 'stal_medium_area' must satisfy 0 < small < medium.") + if not 1.0 <= candidate_scale <= 4.0: + raise ValueError("'stal_candidate_scale' must be between 1.0 and 4.0.") + + integer_values = {"stal_min_candidates": min_candidates, **topks} + for name, value in integer_values.items(): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"'{name}={value}' must be an int.") + if not 1 <= value <= 64: + raise ValueError(f"'{name}={value}' is invalid. Use an integer between 1 and 64.") + if any(value < min_candidates for value in topks.values()): + raise ValueError("Every adaptive STAL top-k must be greater than or equal to 'stal_min_candidates'.") + + def _foundation_transformers_available() -> bool: """Return whether the optional Transformers package can be discovered without importing it.""" try: diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml index ac9aa1572..8283fb695 100644 --- a/ultralytics/cfg/default.yaml +++ b/ultralytics/cfg/default.yaml @@ -107,6 +107,16 @@ warmup_bias_lr: 0.1 # (float) bias learning rate during warmup distill_model: # (str, optional) path to teacher model for knowledge distillation dis: 6.0 # (float) distillation loss weight +# Small-target adaptive label assignment (training-only; default preserves the existing fixed-stride behavior) ------- +stal_mode: fixed # (str) candidate policy: tal (no expansion), fixed (legacy stride expansion), or adaptive +stal_small_area: 1024.0 # (float) small-object threshold in assigner-input pixels (32^2) +stal_medium_area: 9216.0 # (float) medium-object threshold in assigner-input pixels (96^2) +stal_candidate_scale: 1.5 # (float) adaptive expansion factor applied to small-object width and height +stal_min_candidates: 3 # (int) minimum pre-conflict candidate anchors guaranteed for each valid small GT +stal_topk_small: 13 # (int) adaptive top-k for small GT +stal_topk_medium: 10 # (int) adaptive top-k for medium GT +stal_topk_large: 10 # (int) adaptive top-k for large GT + # Foundation teacher distillation (training-only, opt-in; does not affect deployment) ------------------------------- foundation_enabled: False # (bool) enable Foundation Teacher distillation during training foundation_teacher: none # (str) teacher family: none, dinov3, siglip2, or multi diff --git a/ultralytics/utils/loss.py b/ultralytics/utils/loss.py index da08e7be0..adf921d3c 100644 --- a/ultralytics/utils/loss.py +++ b/ultralytics/utils/loss.py @@ -375,6 +375,14 @@ def __init__( beta=6.0, stride=self.stride.tolist(), topk2=tal_topk2, + stal_mode=getattr(h, "stal_mode", "fixed"), + stal_small_area=getattr(h, "stal_small_area", 32**2), + stal_medium_area=getattr(h, "stal_medium_area", 96**2), + stal_candidate_scale=getattr(h, "stal_candidate_scale", 1.5), + stal_min_candidates=getattr(h, "stal_min_candidates", 3), + stal_topk_small=getattr(h, "stal_topk_small", 13), + stal_topk_medium=getattr(h, "stal_topk_medium", 10), + stal_topk_large=getattr(h, "stal_topk_large", 10), ) self.bbox_loss = BboxLoss(m.reg_max).to(device) self.proj = torch.arange(m.reg_max, dtype=torch.float, device=device) diff --git a/ultralytics/utils/tal.py b/ultralytics/utils/tal.py index 630d1d548..fd327b95f 100644 --- a/ultralytics/utils/tal.py +++ b/ultralytics/utils/tal.py @@ -25,6 +25,7 @@ class TaskAlignedAssigner(nn.Module): beta (float): The beta parameter for the localization component of the task-aligned metric. stride (list): List of stride values for different feature levels. stride_val (int): The stride value used for select_candidates_in_gts. + stal_mode (str): Candidate policy: raw TAL, legacy fixed-stride expansion, or adaptive STAL. eps (float): A small value to prevent division by zero. """ @@ -37,6 +38,14 @@ def __init__( stride: list | None = None, eps: float = 1e-9, topk2=None, + stal_mode: str = "fixed", + stal_small_area: float = 32**2, + stal_medium_area: float = 96**2, + stal_candidate_scale: float = 1.5, + stal_min_candidates: int = 3, + stal_topk_small: int = 13, + stal_topk_medium: int = 10, + stal_topk_large: int = 10, ): """Initialize a TaskAlignedAssigner object with customizable hyperparameters. @@ -48,6 +57,14 @@ def __init__( stride (list, optional): List of stride values for different feature levels. eps (float, optional): A small value to prevent division by zero. topk2 (int, optional): Secondary topk value for additional filtering. + stal_mode (str, optional): Candidate policy: ``tal``, ``fixed``, or ``adaptive``. + stal_small_area (float, optional): Small-object threshold in assigner-input pixels. + stal_medium_area (float, optional): Medium-object threshold in assigner-input pixels. + stal_candidate_scale (float, optional): Small-object candidate-box expansion factor. + stal_min_candidates (int, optional): Minimum pre-conflict candidates for each valid small GT. + stal_topk_small (int, optional): Adaptive top-k for small GT. + stal_topk_medium (int, optional): Adaptive top-k for medium GT. + stal_topk_large (int, optional): Adaptive top-k for large GT. """ super().__init__() self.topk = topk @@ -57,6 +74,12 @@ def __init__( self.beta = beta self.stride = stride if stride is not None else [8, 16, 32] self.stride_val = self.stride[1] if len(self.stride) > 1 else self.stride[0] + self.stal_mode = stal_mode + self.stal_small_area = stal_small_area + self.stal_medium_area = stal_medium_area + self.stal_candidate_scale = stal_candidate_scale + self.stal_min_candidates = stal_min_candidates + self.stal_topk = (stal_topk_small, stal_topk_medium, stal_topk_large) self.eps = eps @torch.no_grad() @@ -163,7 +186,10 @@ def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, m # Get anchor_align metric, (b, max_num_obj, h*w) align_metric, overlaps = self.get_box_metrics(pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_in_gts * mask_gt) # Get topk_metric mask, (b, max_num_obj, h*w) - mask_topk = self.select_topk_candidates(align_metric, topk_mask=mask_gt.expand(-1, -1, self.topk).bool()) + if self.stal_mode == "adaptive": + mask_topk = self.select_adaptive_topk_candidates(align_metric, mask_in_gts, gt_bboxes, mask_gt) + else: + mask_topk = self.select_topk_candidates(align_metric, topk_mask=mask_gt.expand(-1, -1, self.topk).bool()) # Merge all mask to a final mask, (b, max_num_obj, h*w) mask_pos = mask_topk * mask_in_gts * mask_gt @@ -246,6 +272,29 @@ def select_topk_candidates(self, metrics, topk_mask=None): return count_tensor.to(metrics.dtype) + def get_adaptive_topks(self, gt_bboxes): + """Return per-GT top-k values from augmented, resized boxes entering the assigner.""" + wh = (gt_bboxes[..., 2:] - gt_bboxes[..., :2]).clamp_(min=0) + areas = wh.prod(-1) + small_topk, medium_topk, large_topk = self.stal_topk + return torch.where( + areas < self.stal_small_area, + small_topk, + torch.where(areas < self.stal_medium_area, medium_topk, large_topk), + ).long() + + def select_adaptive_topk_candidates(self, metrics, candidate_mask, gt_bboxes, mask_gt): + """Select a scale-dependent number of anchors strictly from each GT candidate region.""" + topks = self.get_adaptive_topks(gt_bboxes) + max_topk = min(max(self.stal_topk), metrics.shape[-1]) + masked_metrics = metrics.masked_fill(~candidate_mask.bool(), -torch.inf) + topk_metrics, topk_idxs = torch.topk(masked_metrics, max_topk, dim=-1, largest=True) + rank_mask = torch.arange(max_topk, device=metrics.device).view(1, 1, -1) < topks.unsqueeze(-1) + valid = rank_mask & topk_metrics.isfinite() & mask_gt.bool() + selected = torch.zeros_like(metrics, dtype=torch.int8) + selected.scatter_add_(-1, topk_idxs, valid.to(torch.int8)) + return selected.clamp_(max=1).to(metrics.dtype) + def get_targets(self, gt_labels, gt_bboxes, target_gt_idx, fg_mask): """Compute target labels, target bounding boxes, and target scores for the positive anchor points. @@ -302,17 +351,45 @@ def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt, eps=1e-9): - b: batch size, n_boxes: number of ground truth boxes, h: height, w: width. - Bounding box format: [x_min, y_min, x_max, y_max]. """ - gt_bboxes_xywh = xyxy2xywh(gt_bboxes) - wh_mask = gt_bboxes_xywh[..., 2:] < self.stride[0] # the smallest stride - gt_bboxes_xywh[..., 2:] = torch.where( - (wh_mask * mask_gt).bool(), - torch.tensor(self.stride_val, dtype=gt_bboxes_xywh.dtype, device=gt_bboxes_xywh.device), - gt_bboxes_xywh[..., 2:], - ) - gt_bboxes = xywh2xyxy(gt_bboxes_xywh) + original_bboxes = gt_bboxes + if self.stal_mode == "fixed": + gt_bboxes_xywh = xyxy2xywh(gt_bboxes) + wh_mask = gt_bboxes_xywh[..., 2:] < self.stride[0] # the smallest stride + gt_bboxes_xywh[..., 2:] = torch.where( + (wh_mask * mask_gt).bool(), + torch.tensor(self.stride_val, dtype=gt_bboxes_xywh.dtype, device=gt_bboxes_xywh.device), + gt_bboxes_xywh[..., 2:], + ) + gt_bboxes = xywh2xyxy(gt_bboxes_xywh) + elif self.stal_mode == "adaptive": + gt_bboxes_xywh = xyxy2xywh(gt_bboxes) + areas = gt_bboxes_xywh[..., 2:].prod(-1, keepdim=True) + small_mask = (areas < self.stal_small_area) & mask_gt.bool() + min_extent = 2 * self.stride[0] + expanded_wh = torch.maximum( + gt_bboxes_xywh[..., 2:] * self.stal_candidate_scale, + torch.as_tensor(min_extent, dtype=gt_bboxes.dtype, device=gt_bboxes.device), + ) + gt_bboxes_xywh[..., 2:] = torch.where(small_mask, expanded_wh, gt_bboxes_xywh[..., 2:]) + gt_bboxes = xywh2xyxy(gt_bboxes_xywh) lt, rb = gt_bboxes.unsqueeze(2).chunk(2, 3) # (b, n_boxes, 1, 2) left-top, right-bottom - return ((xy_centers - lt > eps) & (rb - xy_centers > eps)).all(3) + candidate_mask = ((xy_centers - lt > eps) & (rb - xy_centers > eps)).all(3) + if self.stal_mode != "adaptive": + return candidate_mask + + original_wh = (original_bboxes[..., 2:] - original_bboxes[..., :2]).clamp_(min=0) + small_mask = (original_wh.prod(-1) < self.stal_small_area) & mask_gt.squeeze(-1).bool() + needs_candidates = small_mask & (candidate_mask.sum(-1) < self.stal_min_candidates) + if needs_candidates.any(): + centers = (original_bboxes[..., :2] + original_bboxes[..., 2:]) / 2 + distances = (xy_centers.view(1, 1, -1, 2) - centers.unsqueeze(2)).square().sum(-1) + distances.masked_fill_(~needs_candidates.unsqueeze(-1), torch.inf) + nearest = distances.topk(min(self.stal_min_candidates, xy_centers.shape[0]), dim=-1, largest=False).indices + supplements = torch.zeros_like(candidate_mask) + supplements.scatter_(-1, nearest, needs_candidates.unsqueeze(-1).expand_as(nearest)) + candidate_mask |= supplements + return candidate_mask def select_highest_overlaps(self, mask_pos, overlaps, n_max_boxes, align_metric): """Select anchor boxes with highest IoU when assigned to multiple ground truths. @@ -343,10 +420,13 @@ def select_highest_overlaps(self, mask_pos, overlaps, n_max_boxes, align_metric) if self.topk2 != self.topk: align_metric = align_metric * mask_pos # update overlaps + topk2 = min(self.topk2, align_metric.shape[-1]) + if self.stal_mode == "adaptive": + align_metric = align_metric.masked_fill(~mask_pos.bool(), -torch.inf) # (b, n_max_boxes, topk2) - max_overlaps_idx = torch.topk(align_metric, self.topk2, dim=-1, largest=True).indices + topk2_metrics, max_overlaps_idx = torch.topk(align_metric, topk2, dim=-1, largest=True) topk_idx = torch.zeros(mask_pos.shape, dtype=mask_pos.dtype, device=mask_pos.device) # update mask_pos - topk_idx.scatter_(-1, max_overlaps_idx, 1.0) + topk_idx.scatter_(-1, max_overlaps_idx, topk2_metrics.isfinite().to(mask_pos.dtype)) mask_pos *= topk_idx fg_mask = mask_pos.sum(-2) # Find each grid serve which gt(index)