diff --git a/scripts/stal/EXPERIMENTS.md b/scripts/stal/EXPERIMENTS.md new file mode 100644 index 000000000..a648c73b6 --- /dev/null +++ b/scripts/stal/EXPERIMENTS.md @@ -0,0 +1,198 @@ +# A2 STAL experiment report + +This report separates formal accuracy evidence from parameter screening. VisDrone DET does not publish +small/medium/large area metrics. Every `APs`, `APm`, `APl`, and `ARs500` value below is therefore a project +COCO-style supplement computed from original validation-image ground-truth boxes: small `<32^2`, medium +`[32^2, 96^2)`, large `>=96^2`, IoU `.50:.05:.95`, and `maxDets=500`. Training-time STAL gates and assignment +statistics instead use boxes after augmentation and resize on the actual training canvas. + +## Acceptance status + +| Stage | Requested evidence | Status | +|---|---|---| +| P0 | Full VisDrone baseline, scale metrics, and per-epoch positive-assignment statistics | Completed | +| P1 | Adaptive STAL improves formal three-seed APs over pure TAL by at least 1.0 absolute point | Not met | +| P2 | Scan area/scale controls and warmup, or extend to another dataset/task | Completed through the parameter-sensitivity route; no cross-dataset claim | + +P2 completion describes the requested ablation work. It does not override the unmet P1 accuracy threshold. + +## Relation to other A2 submissions + +At review time, #274 emphasizes candidate expansion, per-tier top-k, a minimum-positive path, and detailed evidence +boundaries; #280 reports a compact three-way table and positive-coverage statistics; #293 keeps an A16 +minimum-candidate design and reports its best single-seed attempt. Their training recipes, baselines, assignment rules, +and evaluation handling differ, so their APs values are not merged into the tables below. This submission's distinct +evidence is a matched FP32 three-seed pure/fixed/adaptive comparison plus the parameter and failure analysis in this +document. + +## Formal protocol and result + +All formal rows use complete VisDrone train/val (`6471/548`), YOLO-Master v0.1-N, `imgsz=800`, 120 epochs, +`patience=0`, batch 6, FP32, MuSGD, optimizer warmup 3, Mosaic 1 with `close_mosaic=10`, and the checkpoint selected +by the repository's normal overall-fitness rule. The submitted adaptive arm uses area threshold `0.0016`, constant +relaxation 8, assignment warmup 0, and small-target top-k 10. + +| Assignment | Seed 0 APs (%) | Seed 1 APs (%) | Seed 2 APs (%) | Mean +/- SD (%) | +|---|---:|---:|---:|---:| +| pure TAL | 12.947795 | 13.262371 | 12.880747 | 13.030304 +/- 0.203752 | +| fixed-stride STAL | 13.064431 | 12.961409 | 13.156323 | 13.060721 +/- 0.097510 | +| adaptive-w0 | 13.414071 | 13.319065 | 13.135239 | 13.289458 +/- 0.141754 | + +The paired adaptive-minus-pure mean is `+0.259154` absolute point with a 95% Student-t interval of +`[-0.249674, +0.767982]`. Relative to fixed STAL it is `+0.228738` point with interval +`[-0.308802, +0.766278]`. Both intervals cross zero, and the mean gain is below 1.0 point. The Mosaic-off seed-0 +pair is also positive but small: pure `11.692718%`, adaptive `12.001212%`, delta `+0.308493` point. +The paired statistics use the unrounded evaluator outputs; the per-seed values displayed in the table are rounded to +six decimal places. + +The retained prediction artifacts also support the required secondary scale metrics. The values below were +recomputed from the official-format TXT files, whose box coordinates are rounded to three decimals; this changes APs +by at most `0.0126` absolute point relative to the unrounded JSON values above. They are reported here for the +medium/large and small-recall context, while the unrounded JSON table remains the P1 acceptance source. + +| Assignment | Seed | APm (%) | APl (%) | ARs500 (%) | +|---|---:|---:|---:|---:| +| pure TAL | 0 | 32.196537 | 41.687356 | 29.387689 | +| pure TAL | 1 | 32.774068 | 40.917719 | 29.232980 | +| pure TAL | 2 | 32.082815 | 44.245727 | 29.016809 | +| fixed-stride STAL | 0 | 31.888752 | 41.420740 | 29.430121 | +| fixed-stride STAL | 1 | 32.064242 | 44.290912 | 29.419533 | +| fixed-stride STAL | 2 | 32.016438 | 42.538901 | 29.394237 | +| adaptive-w0 | 0 | 31.756169 | 40.036916 | 29.287789 | + +The official `VisDrone2018-DET-toolkit` commit `005445782213e20cb91bc50a597db3dd949e749a` produced the +following overall metrics on all 548 original validation annotations. These official metrics do not contain area +bins and are not substituted for the project APs above. + +| Assignment | Seed | Official AP (%) | AP50 (%) | AP75 (%) | AR500 (%) | +|---|---:|---:|---:|---:|---:| +| pure TAL | 0 | 22.450013 | 40.160397 | 21.507844 | 39.137572 | +| pure TAL | 1 | 22.602562 | 40.447796 | 21.710440 | 39.197234 | +| pure TAL | 2 | 22.404797 | 40.309586 | 21.363486 | 38.831045 | +| fixed-stride STAL | 0 | 22.636834 | 40.990155 | 21.477616 | 39.205206 | +| fixed-stride STAL | 1 | 22.921176 | 41.358936 | 21.750081 | 39.233329 | +| fixed-stride STAL | 2 | 22.741519 | 41.152091 | 21.607634 | 39.226400 | +| adaptive-w0 | 0 | 22.743838 | 41.140512 | 21.415664 | 39.090666 | + +The adaptive-w0 seed-1/2 predictions have project APs results but are absent from the retained official-evaluation +package, so no official three-seed adaptive mean is claimed. + +## P2 parameter sensitivity + +These screens choose mechanisms and expose tradeoffs; they are not formal P1 results. Each table states its own +budget so that short-run values are not mixed with the formal table. + +### Area gate and relaxation + +The original relative-area gate was 1%. Screens narrowed the affected training targets to 0.5%, 0.3%, and 0.16%; the +0.16% gate was the only candidate promoted to a full 120-epoch check. At seed 0, moving from the original adaptive +arm to the 0.16% gate with relaxation 8 changed APs from `13.167447%` to `13.337117%`. A mild area-adaptive top-k +variant reached `13.364341%`, only `+0.027224` point beyond the fixed top-k version. This supports narrowing the +intervention, but the final three-seed gain remained small. + +A separate 10%-train, 10-epoch screen varied only constant relaxation. Its endpoint is framework mAP50-95, not APs: + +| Total width/height relaxation (px) | mAP50-95 | Small positives / GT | Small zero-positive ratio | +|---:|---:|---:|---:| +| 4 | 0.01176 | 3.71981 | 0.153254 | +| 6 | 0.00891 | 3.97614 | 0.159743 | +| 8 | 0.01103 | 4.17120 | 0.156968 | + +More expansion increased positive count, but it did not monotonically improve the zero-positive ratio or accuracy. +On fixed real batches, increasing relaxation from 0 to 8 also raised conflict loss among pre-conflict-covered targets +from `10.49%` to `12.31%`, while the fraction with alignment above epsilon stayed near `7.64%`. + +Continuous square-root area scaling was also negative on the same 10%-train, 10-epoch protocol: + +| Relaxation rule | APs | +|---|---:| +| constant r8 control | 0.009286 | +| sqrt-area r0-to-r8 | 0.007113 | +| sqrt-area r4-to-r8 | 0.006761 | + +### Warmup and top-k + +The full-data 24-epoch prefix screen retained the 120-epoch learning-rate schedule and evaluated fixed checkpoints. + +| Assignment warmup | Small top-k | Epoch-24 APs (%) | Epoch-24 ARs500 (%) | +|---:|---:|---:|---:| +| 0 | 10 | 9.207494 | 23.019632 | +| 5 | 10 | 8.964761 | 23.392646 | +| 10 | 10 | 8.626137 | 22.921022 | +| 10 | 5 | 8.685307 | 23.082611 | + +Warmup 0 led at the fixed 24-epoch checkpoint and was therefore promoted. This is a screening decision rather than +evidence that 24-epoch ranking predicts 120-epoch ranking. A retrospective four-arm check found Spearman correlation +between early framework mAP50-95 rank and final APs rank of `-0.8` at epoch 10, `0.0` at epoch 20, and `0.8` at epoch +40; four arms are enough to reject epoch-10 ranking as a dependable gate, not enough to establish a general predictor. + +An independent area-adaptive top-k screen used 10% train for 10 epochs: + +| Tiny-to-threshold top-k | mAP50-95 | Small positives / GT | Target score / positive | +|---|---:|---:|---:| +| 10-to-10 control | 0.01103 | 4.17120 | 0.153921 | +| 2-to-10 | 0.01091 | 2.26906 | 0.252716 | +| 5-to-10 | 0.01058 | 3.29717 | 0.187562 | + +Reducing the nomination budget improved average selected-positive quality but reduced recall and did not improve the +screen endpoint. Uniform larger top-k arms and the final 8-to-10 variant likewise did not show a material advantage. + +### Coverage, supervision quality, and auxiliary loss + +The final epoch of a matched seed-0 telemetry comparison reports all three training-canvas size bins. The focus row +uses the selected `0.0016` area gate and constant relaxation 8; it predates the later warmup-0 selection, so it is +mechanism evidence rather than a substitute for the adaptive-w0 accuracy table. + +| Assignment | Small pos/GT | Small zero (%) | Medium pos/GT | Medium zero (%) | Large pos/GT | Large zero (%) | +|---|---:|---:|---:|---:|---:|---:| +| pure TAL | 3.50122 | 18.8060 | 9.93647 | 0.024975 | 9.98187 | 0.000000 | +| fixed-stride STAL | 4.05073 | 6.92296 | 9.93314 | 0.038597 | 9.98302 | 0.000000 | +| focus t0016-r8 | 5.79637 | 6.26367 | 9.93364 | 0.027245 | 9.97994 | 0.000000 | + +A 1618-image, 24-epoch quality-aware coverage screen reduced the training zero-positive ratio from `10.2975%` to +`8.9193%` and raised ARs500 by `0.2018` point, while APs changed from `4.127463%` to `4.108997%`. Likewise, adding at +most one or two adaptive-only candidates sharply reduced zero-positive incidence in an earlier screen but lowered +APs. These paired results directly show that positive coverage alone is not a sufficient promotion criterion. + +Changing MoE auxiliary-loss strength was kept separate from STAL geometry. Under auxiliary strength 1, focus-minus- +pure APs was `+0.389322` point at seed 0; under strength 3 it was only `+0.075472` point. The interaction is evidence +that optimization context matters, but auxiliary strength is not claimed as an STAL contribution. + +## Why the gain stopped below one point + +The following are confirmed observations: + +- Candidate expansion activates and materially changes assignment. In the original formal seed-0 comparison it + raised augmented-small positives per GT from `4.045` (fixed) to `5.904` and reduced zero-positive incidence from + `6.98%` to `5.95%`. +- The extra assignments are not monotonic evidence of useful supervision. Relaxation increased conflicts, and + coverage-focused screens improved zero-positive or recall statistics while APs stayed flat or fell. +- Fixed versus original adaptive prediction analysis showed a larger gain at AP50s (`+0.3761` point) but a small loss + at AP75s (`-0.0489` point) and essentially unchanged ARs500 (`-0.0153` point). At confidence 0.1, unmatched + small-prediction diagnostics counted 1356 more duplicate/competition boxes and 1108 more localization/mixed-error + boxes for adaptive. These are geometric diagnostic labels, not an additive AP-loss decomposition. +- Per-class APs moved in both directions: car improved by `+0.577` point, while people, van, and truck changed by + `-0.242`, `-0.122`, and `-0.251` point in that seed-0 comparison. +- The final three-seed gain varies enough that its confidence interval crosses zero. + +The most plausible interpretation is that candidate expansion solves part of the coverage problem but also adds weak +or competing supervision. That can improve coarse-IoU detections while failing to improve stricter localization and +ranking enough for APs averaged over ten IoU thresholds. The training gate also acts on augmented relative area, +whereas APs is grouped by original-image area, so not every evaluated small object receives the same intervention. +These mechanisms are consistent with the observations but have not been isolated as a complete causal decomposition. + +The saved predictions are post-processing outputs (`conf>=0.001`, at most 500 detections per image), so they cannot +separate failures originating in the raw head from NMS or output truncation. The report therefore does not attribute a +numerical fraction of the AP gap to localization, classification, duplicates, or missed detections. + +## Reproduction and validation boundaries + +Use the commands in [README.md](README.md) for training, supplemental scale evaluation, and official VisDrone TXT +export. The submitted code has focused tests for candidate geometry, configuration contracts, pure/fixed/adaptive +behavior, empty and overlapping targets, conflict handling, FP16 area arithmetic, AMP finite loss/gradients, +telemetry, scale evaluation, and VisDrone export validation. + +The official VisDrone toolkit was run separately on the original 548-image annotation set for pure and fixed across +all three seeds and adaptive seed 0. Official AP/AP50/AP75/AR values do not contain area bins and are not substituted +for project APs. The full P2 route in this report is parameter sensitivity on VisDrone; no DOTA/AI-TOD, segmentation, +or pose generalization result is claimed. diff --git a/scripts/stal/README.md b/scripts/stal/README.md new file mode 100644 index 000000000..4745b8140 --- /dev/null +++ b/scripts/stal/README.md @@ -0,0 +1,160 @@ +# STAL experiment workflow + +This directory contains the diagnostics, evaluation, and VisDrone export utilities for the A2 small-target adaptive +label-assignment experiments. The formal protocol uses three seeds, a separate Mosaic interaction control, and a clear +distinction between official VisDrone metrics and the project's supplemental COCO-style scale metrics. + +The formal results, P2 parameter-sensitivity tables, and evidence-bounded explanation of the unmet P1 target are in +[EXPERIMENTS.md](EXPERIMENTS.md). + +## Training controls + +The locked fixed-stride behavior remains the default because `stal_candidate_mode=fixed`, `stal_enabled=False`, and +`stal_stats=False`. The legacy `stal_enabled=True` switch remains supported and selects adaptive mode. + +| Argument | Meaning | Default | +|---|---|---:| +| `stal_enabled` | Enable relative-area candidate-region relaxation | `False` | +| `stal_stats` | Record assignment counts in `results.csv` | `False` | +| `stal_candidate_mode` | Candidate policy: `pure`, `fixed`, or `adaptive` | `fixed` | +| `stal_small_topk` | Adaptive-mode top-k for small targets; standard TAL remains unchanged | `10` | +| `stal_area_threshold` | Target-to-current-training-image area ratio used by STAL | `0.01` | +| `stal_relaxation` | Maximum total width and height added to the candidate region | `8.0` pixels | +| `stal_warmup_epochs` | Linear ramp duration for the relaxation | `10.0` epochs | +| `stal_zero_positive_rescue` | Rescue an uncovered small GT with its best legal positive-quality candidate | `False` | + +`stal_relaxation=8` means adding four pixels on every side because the implementation operates on center-width-height +boxes and increases total width and height by eight pixels. + +Run the required three-way candidate-policy comparison with statistics enabled: + +```bash +yolo train model= data= stal_candidate_mode=pure stal_stats=True \ + project= name=pure-tal-seed0 seed=0 +yolo train model= data= stal_candidate_mode=fixed stal_stats=True \ + project= name=fixed-stride-seed0 seed=0 +yolo train model= data= stal_candidate_mode=adaptive stal_stats=True \ + stal_area_threshold=0.01 stal_relaxation=8 stal_warmup_epochs=10 \ + project= name=adaptive-stal-seed0 seed=0 +``` + +`pure` uses the unmodified GT candidate region. `fixed` preserves the repository's existing rule that expands a GT +dimension below the smallest stride to the fixed middle stride. `adaptive` keeps that existing behavior and adds the +relative-area relaxation. For backward compatibility, `stal_enabled=True` also selects `adaptive`; it conflicts with +`stal_candidate_mode=pure`. + +For the formal P0/P1 protocol use complete VisDrone train, all 548 validation images, YOLO-Master v0.1-N, `imgsz=800`, +`epochs=120`, and `patience=0`. Subset or short-cycle experiments are screening/smoke evidence only. + +The submitted three-seed `adaptive-w0` arm used `stal_candidate_mode=adaptive`, `stal_area_threshold=0.0016`, +`stal_relaxation=8`, `stal_warmup_epochs=0`, and `stal_small_topk=stal_small_topk_min=10`. Candidate-quality gates, +minimum-candidate rescue, NWD, and SimD remained disabled so the comparison isolates candidate-region relaxation. + +Run adaptive STAL using the backward-compatible flag if needed: + +```bash +yolo train model= data= \ + stal_enabled=True stal_stats=True \ + stal_area_threshold=0.01 stal_relaxation=8 stal_warmup_epochs=10 \ + project= name=stal-seed0 seed=0 +``` + +The following groups are added to `results.csv` when `stal_stats=True`: `stal`, COCO-style `small`/`medium`/`large`, +and `all`. Every group reports GT count, positive count, positives per GT, zero-positive GT count, and zero-positive +ratio. The original compatibility columns remain, including: + +- `assign/stal_gt` and `assign/stal_pos` +- `assign/pos_per_stal_gt` +- `assign/all_gt` and `assign/all_pos` +- `assign/pos_per_gt` + +All training groups are defined on the augmented, resized training canvas. This is mechanism evidence and is +deliberately separate from validation AP, which uses GT bbox area in the original validation images. + +Score diagnostics now distinguish `assign/stal_*` (relative area below `stal_area_threshold`) from +`assign/small_*` (absolute bbox area below 32 squared pixels), alongside `assign/all_*`. Historical +`small_nonzero_score_*` and `small_target_score_*` columns used the relative STAL gate; do not interpret those old +columns as COCO-style small statistics. Start a new output directory with this schema; incompatible CSV appends fail. +E2E models report one-to-many diagnostics and drain both branches, avoiding duplicate GT counts. Counters reset +at every epoch attempt so OOM replay does not include abandoned batches. +The unconditional nearest-candidate guarantee cannot be combined with a positive geometric IoU capacity floor; +otherwise it could reintroduce a candidate rejected by that floor. Both remain disabled by default. + +When `stal_zero_positive_rescue=True`, `results.csv` also records a staged rescue funnel: + +- `assign/rescue_attempted`: uncovered STAL-area GTs presented to rescue. +- `assign/rescue_has_legal_candidate`: attempted GTs with at least one candidate in the relaxed region. +- `assign/rescue_has_free_candidate`: attempted GTs with a legal candidate not already assigned to another GT. +- `assign/rescue_has_positive_quality_candidate`: attempted GTs whose best free candidate has a finite positive + task-alignment score. +- `assign/rescue_proposed`, `assign/rescue_succeeded`, and `assign/rescue_lost_to_conflict`: proposals before the + second conflict pass, surviving rescues, and proposals removed by that conflict pass. +- `assign/rescue_bootstrap_proposed` and `assign/rescue_bootstrap_succeeded`: zero-quality targets proposed and retained + through the optional center-prior bootstrap path. + +These counters diagnose why rescue does or does not reduce zero-positive targets. They are assignment-mechanism +evidence, not an AP improvement claim. + +With `stal_stats=True`, the behavior-neutral assignment funnel also records: + +- `assign/stal_no_legal_candidate`: STAL-area GTs with no point in the candidate region. +- `assign/stal_legal_zero_alignment`: candidate-covered STAL-area GTs whose alignment is exactly zero. +- `assign/stal_sub_eps_alignment`: STAL-area GTs with positive alignment no greater than assigner epsilon. +- `assign/stal_above_eps_alignment`: STAL-area GTs with alignment above assigner epsilon. +- `assign/stal_topk_missed_nonzero`: nonzero-alignment STAL-area GTs not selected before conflict resolution. +- `assign/stal_preconflict_positive`: STAL-area GTs selected before conflict resolution. +- `assign/stal_conflict_lost`: selected STAL-area GTs left uncovered by conflict resolution. +- `assign/stal_postconflict_zero`: all STAL-area GTs left uncovered before optional rescue. + +The funnel is collected only when statistics are enabled and does not modify candidate masks, ranking, target scores, +or losses. + +The bootstrap path is disabled by default. Set `stal_rescue_score_floor` to a value in `(0, 1]` to let an uncovered +small GT with no positive task-alignment score nominate its nearest free legal anchor. The selected anchor receives that +value as a minimum target-score weight, so it contributes bounded classification and localization supervision instead +of a zero-weight label. `stal_rescue_floor_decay_epochs=N` linearly decays this floor to zero by epoch `N`; zero keeps +the configured floor constant. A positive floor requires both `stal_zero_positive_rescue=True` and adaptive mode. + +## Size-binned validation AP + +First emit Ultralytics prediction JSON from the selected checkpoint: + +```bash +yolo val model= data= save_json=True max_det=500 \ + project= name= +``` + +Then build COCO ground truth from the YOLO labels and evaluate the same predictions: + +```bash +python scripts/stal/evaluate_scale_ap.py \ + --data \ + --predictions //predictions.json \ + --annotations-out //visdrone-val-coco.json \ + --metrics-out //scale-ap.json +``` + +The evaluator uses `maxDets=[1,10,500]` and reports `AP`, `AP50`, `AP75`, `APs`, `APm`, `APl`, `AP50s`, `AR500`, +`ARs500`, `ARm500`, and `ARl500`. The bins are the project's COCO-style supplemental definition on original-image +GT bbox area: small `<32²`, medium `[32²,96²)`, and large `>=96²`. They are not official VisDrone area bins. + +The formal P1 gate is an absolute increase of at least `0.01` in `APs@[IoU=.50:.95,maxDets=500]`, i.e. 1.0 AP point. +Use original VisDrone annotations and an official-compatible evaluator for claims labelled official VisDrone AP/AP50/ +AP75/AR500; a ground truth reconstructed from YOLO labels may not retain ignore/truncation/occlusion metadata. +Install `faster-coco-eval>=1.6.7` in the experiment environment before running it. +Evaluation now rejects missing label files; use an explicitly empty label file for a verified background image. +All returned AP/AR values are fractions, with `-1` for unavailable GT bins. The CLI explicitly marks these as +supplemental metrics rather than official VisDrone scores. + +Export the same predictions to the official eight-field VisDrone DET submission layout (including empty result files): + +```bash +python scripts/stal/export_visdrone_results.py \ + --predictions //predictions.json \ + --images /images \ + --output //visdrone-official-results +``` + +Run those TXT files with the official `VisDrone2018-DET-toolkit` and the original annotation directory. The official +toolkit removes detections in ignored regions and reports AP/AP50/AP75 plus AR@1/10/100/500. Keep its console output as +a separate evidence artifact from `scale-ap.json`. diff --git a/scripts/stal/__init__.py b/scripts/stal/__init__.py new file mode 100644 index 000000000..e563a8123 --- /dev/null +++ b/scripts/stal/__init__.py @@ -0,0 +1 @@ +"""Utilities for STAL experiments and evaluation.""" diff --git a/scripts/stal/diagnose_amp_training.py b/scripts/stal/diagnose_amp_training.py new file mode 100644 index 000000000..dae4c371f --- /dev/null +++ b/scripts/stal/diagnose_amp_training.py @@ -0,0 +1,325 @@ +#!/usr/bin/env python3 +"""Compare short FP32 and AMP training from identical YOLO-Master weights.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +import random +from itertools import zip_longest +from pathlib import Path + +import numpy as np +import torch + +from ultralytics import YOLO +from ultralytics.models.yolo.detect import DetectionTrainer + + +def tensor_digest(tensor: torch.Tensor) -> str: + """Hash tensor metadata and all bytes, preserving element order and dtype.""" + value = tensor.detach().cpu().contiguous() + digest = hashlib.sha256(str((str(value.dtype), tuple(value.shape))).encode()) + digest.update(value.reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def state_digest(state: dict[str, torch.Tensor]) -> str: + """Hash a complete named model state including buffers.""" + return hashlib.sha256( + json.dumps({name: tensor_digest(value) for name, value in sorted(state.items())}).encode() + ).hexdigest() + + +def shared_state_trainer(initial_state: dict[str, torch.Tensor]): + """Restore shared weights after dataset-specific model creation, before optimizer and EMA setup.""" + + class SharedStateTrainer(DetectionTrainer): + def get_model(self, cfg=None, weights=None, verbose=True): + model = super().get_model(cfg=cfg, weights=weights, verbose=verbose) + if not initial_state: + initial_state.update(clone_state(model)) + model.load_state_dict(initial_state, strict=True) + return model + + return SharedStateTrainer + + +def batch_signature(batch: dict) -> dict: + """Identify full ordered detection inputs, including target-to-image associations.""" + return { + "files": [str(path) for path in batch.get("im_file", ())], + "tensors": { + key: tensor_digest(batch[key]) if isinstance(batch.get(key), torch.Tensor) else None + for key in ("img", "bboxes", "cls", "batch_idx") + }, + } + + +def first_batch_mismatch(left: list[str], right: list[str]) -> int | None: + """Report differing content or the first missing batch, using zero-based indices.""" + return next((i for i, pair in enumerate(zip_longest(left, right)) if pair[0] != pair[1]), None) + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Model YAML used to create one shared initial state.") + parser.add_argument("--data", required=True, help="Dataset YAML.") + parser.add_argument("--output", type=Path, required=True, help="Output directory for runs and summary JSON.") + parser.add_argument("--device", default="0") + parser.add_argument("--imgsz", type=int, default=800) + parser.add_argument("--batch", type=int, default=4) + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--fraction", type=float, default=0.02) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--candidate-mode", choices=("pure", "fixed", "adaptive"), default="pure") + parser.add_argument("--rescue", action="store_true") + parser.add_argument("--rescue-score-floor", type=float, default=0.0) + parser.add_argument("--rescue-floor-decay-epochs", type=float, default=0.0) + return parser.parse_args() + + +def seed_everything(seed: int) -> None: + """Reset host and CUDA random generators before constructing each arm.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def clone_state(model: torch.nn.Module) -> dict[str, torch.Tensor]: + """Clone a model state onto CPU for identical-arm restoration.""" + return {name: value.detach().cpu().clone() for name, value in model.state_dict().items()} + + +def clone_parameters(model: torch.nn.Module) -> dict[str, torch.Tensor]: + """Clone trainable parameters onto CPU for update-delta diagnostics.""" + return {name: value.detach().cpu().float().clone() for name, value in model.named_parameters()} + + +def parameter_delta(initial: dict[str, torch.Tensor], model: torch.nn.Module) -> dict[str, float | int]: + """Summarize parameter movement relative to a shared initial state.""" + delta_sq = 0.0 + initial_sq = 0.0 + max_abs = 0.0 + changed = 0 + total = 0 + for name, parameter in model.named_parameters(): + before = initial[name] + after = parameter.detach().cpu().float() + delta = after - before + delta_sq += float(delta.square().sum()) + initial_sq += float(before.square().sum()) + max_abs = max(max_abs, float(delta.abs().max())) + changed += int(bool(torch.count_nonzero(delta))) + total += 1 + delta_l2 = math.sqrt(delta_sq) + initial_l2 = math.sqrt(initial_sq) + return { + "parameter_tensors": total, + "changed_parameter_tensors": changed, + "parameter_delta_l2": delta_l2, + "parameter_relative_delta_l2": delta_l2 / max(initial_l2, 1e-12), + "parameter_max_abs_delta": max_abs, + } + + +def last_csv_row(path: Path) -> dict[str, str]: + """Read the last metrics row when a diagnostic arm emitted one.""" + if not path.exists(): + return {} + with path.open(encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + return rows[-1] if rows else {} + + +def run_arm(args: argparse.Namespace, initial_state: dict[str, torch.Tensor], amp: bool) -> dict: + """Run one diagnostic arm and collect optimizer/scaler/parameter evidence.""" + seed_everything(args.seed) + yolo = YOLO(args.model) + initial_parameters = {} + arm = "amp" if amp else "fp32" + trace = { + "arm": arm, + "batches": 0, + "loss_finite": True, + "loss_items_finite": True, + "gradient_nonfinite_seen": False, + "scaler_scales": [], + "scale_decreases": 0, + "scale_increases": 0, + "optimizer_steps": 0, + "epoch_passes": [], + } + + def on_train_start(trainer) -> None: + trainer.final_eval = lambda: None + trace["initial_state_sha256"] = state_digest(trainer.model.state_dict()) + if trace["initial_state_sha256"] != state_digest(initial_state): + raise RuntimeError("Training model no longer matches the shared initial state") + initial_parameters.update(clone_parameters(trainer.model)) + # AMP capability checks may consume random numbers before the dataloader iterator is created. + # Reset at the shared training boundary so both arms receive identical augmented batches. + seed_everything(args.seed) + scale = float(trainer.scaler.get_scale()) + trace["initial_scaler_scale"] = scale + trace["scaler_scales"].append(scale) + + def on_train_epoch_start(trainer) -> None: + trace["epoch_passes"].append( + { + "epoch": int(trainer.epoch), + "amp_enabled": bool(trainer.amp), + "scaler_scale_start": float(trainer.scaler.get_scale()), + "optimizer_steps_start": int(getattr(trainer, "optimizer_steps", 0)), + "gradient_nonfinite_seen": False, + "_batch_signatures": [], + } + ) + + def on_train_batch_end(trainer) -> None: + trace["batches"] += 1 + trace["optimizer_steps"] = int(getattr(trainer, "optimizer_steps", 0)) + trace["gradient_nonfinite_seen"] |= bool(getattr(trainer, "_gradient_nonfinite", False)) + current_pass = trace["epoch_passes"][-1] + current_pass["gradient_nonfinite_seen"] |= bool(getattr(trainer, "_gradient_nonfinite", False)) + diagnostic = getattr(trainer, "_nonfinite_diagnostic", None) + if diagnostic is not None and diagnostic not in current_pass.setdefault("nonfinite_diagnostics", []): + current_pass["nonfinite_diagnostics"].append(dict(diagnostic)) + current_pass["optimizer_steps_end"] = int(getattr(trainer, "optimizer_steps", 0)) + current_pass["scaler_scale_end"] = float(trainer.scaler.get_scale()) + loss = getattr(trainer, "loss", None) + items = getattr(trainer, "loss_items", None) + if isinstance(loss, torch.Tensor): + trace["loss_finite"] &= bool(torch.isfinite(loss.detach()).all().item()) + if isinstance(items, torch.Tensor): + trace["loss_items_finite"] &= bool(torch.isfinite(items.detach()).all().item()) + trace["last_loss_items"] = [float(value) for value in items.detach().float().cpu().reshape(-1)] + scale = float(trainer.scaler.get_scale()) + previous = trace["scaler_scales"][-1] + trace["scale_decreases"] += int(scale < previous) + trace["scale_increases"] += int(scale > previous) + trace["scaler_scales"].append(scale) + + def on_train_batch_start(trainer) -> None: + trace["epoch_passes"][-1]["_batch_signatures"].append(batch_signature(trainer.batch)) + + yolo.add_callback("on_train_start", on_train_start) + yolo.add_callback("on_train_epoch_start", on_train_epoch_start) + yolo.add_callback("on_train_batch_start", on_train_batch_start) + yolo.add_callback("on_train_batch_end", on_train_batch_end) + run_name = f"{args.candidate_mode}-{arm}-fraction-{str(args.fraction).replace('.', 'p')}" + yolo.train( + trainer=shared_state_trainer(initial_state), + data=args.data, + epochs=1, + imgsz=args.imgsz, + batch=args.batch, + device=args.device, + workers=args.workers, + pretrained=False, + optimizer="auto", + seed=args.seed, + deterministic=True, + patience=0, + amp=amp, + mosaic=0.0, + plots=False, + save=False, + val=False, + fraction=args.fraction, + verbose=False, + project=str(args.output), + name=run_name, + exist_ok=True, + stal_candidate_mode=args.candidate_mode, + stal_stats=True, + stal_enabled=False, + stal_zero_positive_rescue=args.rescue, + stal_rescue_score_floor=args.rescue_score_floor, + stal_rescue_floor_decay_epochs=args.rescue_floor_decay_epochs, + ) + trace["final_scaler_scale"] = float(yolo.trainer.scaler.get_scale()) + if not initial_parameters: + raise RuntimeError("Training did not reach on_train_start; parameter deltas are unavailable.") + trace.update(parameter_delta(initial_parameters, yolo.model)) + metrics_path = Path(yolo.trainer.save_dir) / "results.csv" + trace["results_csv"] = str(metrics_path) + trace["metrics"] = last_csv_row(metrics_path) + trace["scaler_scales"] = sorted(set(trace["scaler_scales"])) + for epoch_pass in trace["epoch_passes"]: + batch_signatures = epoch_pass.pop("_batch_signatures") + signature_payloads = [ + json.dumps(signature, sort_keys=True, separators=(",", ":")).encode() for signature in batch_signatures + ] + epoch_pass["batches"] = len(batch_signatures) + epoch_pass["batch_signature_sha256"] = hashlib.sha256(b"\n".join(signature_payloads)).hexdigest() + epoch_pass["batch_signature_items_sha256"] = [ + hashlib.sha256(payload).hexdigest() for payload in signature_payloads + ] + epoch_pass["batch_signature_preview"] = batch_signatures[:3] + primary_pass = trace["epoch_passes"][0] + trace["batch_signature_sha256"] = primary_pass["batch_signature_sha256"] + trace["batch_signature_items_sha256"] = primary_pass["batch_signature_items_sha256"] + trace["batch_signature_preview"] = primary_pass["batch_signature_preview"] + trace["epoch_pass_count"] = len(trace["epoch_passes"]) + trace["precision_fallback_seen"] = ( + amp and primary_pass["amp_enabled"] and any(not item["amp_enabled"] for item in trace["epoch_passes"][1:]) + ) + del yolo + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return trace + + +def main() -> int: + """Run both precision arms and emit a machine-readable diagnostic report.""" + args = parse_args() + if not 0 < args.fraction <= 1: + raise ValueError("--fraction must be in (0, 1].") + if args.rescue and args.candidate_mode != "adaptive": + raise ValueError("--rescue requires --candidate-mode adaptive.") + args.output.mkdir(parents=True, exist_ok=True) + seed_everything(args.seed) + initial_state = {} + arms = [run_arm(args, initial_state, amp=False), run_arm(args, initial_state, amp=True)] + report = { + "protocol": { + "model": args.model, + "data": args.data, + "device": args.device, + "imgsz": args.imgsz, + "batch": args.batch, + "fraction": args.fraction, + "seed": args.seed, + "candidate_mode": args.candidate_mode, + "rescue": args.rescue, + "rescue_score_floor": args.rescue_score_floor, + "rescue_floor_decay_epochs": args.rescue_floor_decay_epochs, + "shared_initial_state": arms[0]["initial_state_sha256"] == arms[1]["initial_state_sha256"], + }, + "arms": arms, + "comparison": { + "identical_batch_signatures": arms[0]["batch_signature_sha256"] == arms[1]["batch_signature_sha256"], + "first_mismatched_batch": first_batch_mismatch( + arms[0]["batch_signature_items_sha256"], arms[1]["batch_signature_items_sha256"] + ), + "both_losses_finite": all(arm["loss_finite"] and arm["loss_items_finite"] for arm in arms), + "any_nonfinite_gradient": any(arm["gradient_nonfinite_seen"] for arm in arms), + "amp_precision_fallback_seen": arms[1]["precision_fallback_seen"], + }, + } + report_path = args.output / "amp-training-diagnostic.json" + report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(json.dumps(report, indent=2)) + print(f"REPORT: {report_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/stal/diagnose_rescue_candidates.py b/scripts/stal/diagnose_rescue_candidates.py new file mode 100644 index 000000000..55a98e185 --- /dev/null +++ b/scripts/stal/diagnose_rescue_candidates.py @@ -0,0 +1,191 @@ +"""Observe unused candidate quality during a short, unchanged pure TAL training run.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from collections import Counter +from pathlib import Path +from typing import ClassVar + +import torch + +from ultralytics import YOLO +from ultralytics.utils.tal import TaskAlignedAssigner + + +def candidate_counts(mask_pos, quality, metric, legal, small, eps): + """Count uncovered GTs with free candidates; never change assignment inputs.""" + uncovered = small & ~mask_pos.bool().any(-1) + free = ~mask_pos.bool().any(-2) + free_legal = legal.bool() & free.unsqueeze(1) + has_free = free_legal.any(-1) + best_quality = quality.masked_fill(~free_legal, -torch.inf).amax(-1) + best_metric = metric.masked_fill(~free_legal, -torch.inf).amax(-1) + eligible = uncovered & has_free + suppressed = eligible & (best_metric <= eps) + counts = { + "small_gt": small.sum(), + "uncovered": uncovered.sum(), + "uncovered_without_legal": (uncovered & ~legal.bool().any(-1)).sum(), + "uncovered_with_free_legal": eligible.sum(), + "free_positive_ciou": (eligible & (best_quality > 0)).sum(), + "free_alignment_above_eps": (eligible & (best_metric > eps)).sum(), + "positive_ciou_but_alignment_below_eps": (suppressed & (best_quality > 0)).sum(), + } + for threshold in (0.01, 0.03, 0.05, 0.1, 0.2): + counts[f"suppressed_ciou_ge_{threshold}"] = (suppressed & (best_quality >= threshold)).sum() + return {key: int(value.item()) for key, value in counts.items()} + + +def outside_capacity_counts(boxes, anchors, without_legal, foreground): + """Bound IoU for boxes decoded from nonnegative distances at existing feature points. + + The tightest enclosing rectangle of a GT and a point maximizes IoU over rectangles + containing that point. This ignores DFL's finite range, so it is an optimistic bound. + """ + counts = Counter() + for batch_idx in range(boxes.shape[0]): + selected = boxes[batch_idx, without_legal[batch_idx]] + free = ~foreground[batch_idx].bool() + for chunk in selected.split(128): + if not chunk.numel(): + continue + lower = torch.minimum(chunk[:, None, :2], anchors[None]) + upper = torch.maximum(chunk[:, None, 2:], anchors[None]) + area = (chunk[:, 2:] - chunk[:, :2]).prod(-1) + bound = area[:, None] / (upper - lower).prod(-1).clamp_min(1e-12) + best = bound.amax(-1) + best_free = bound.masked_fill(~free[None], 0).amax(-1) + counts["outside_gt"] += len(chunk) + for threshold in (0.5, 0.75, 0.95): + counts[f"outside_any_iou_cap_ge_{threshold}"] += int((best >= threshold).sum()) + counts[f"outside_free_iou_cap_ge_{threshold}"] += int((best_free >= threshold).sum()) + return dict(counts) + + +class DiagnosticAssigner(TaskAlignedAssigner): + """Collect counts without changing masks, target scores, or optimization.""" + + records: ClassVar[list] = [] + active_epoch = None + + def get_pos_mask(self, *args, **kwargs): + """Retain candidate tensors only until the conflict-resolution call.""" + result = super().get_pos_mask(*args, **kwargs) + self._diagnostic_context = None + if self.active_epoch is not None: + if self.stal_candidate_mode != "pure" or self.stal_nwd_weight or self.stal_zero_positive_rescue: + raise ValueError("Candidate diagnostics require pure TAL with NWD and rescue disabled.") + boxes, valid = args[3], args[5] + image_size = torch.as_tensor(kwargs["image_size"], device=boxes.device, dtype=boxes.dtype) + small = self.small_target_mask(boxes, image_size) & valid.squeeze(-1).bool() + self._diagnostic_context = (result[3], small, boxes, args[4]) + return result + + def select_highest_overlaps(self, mask_pos, overlaps, n_max_boxes, align_metric): + """Observe final pure assignments before target-score normalization.""" + result = super().select_highest_overlaps(mask_pos, overlaps, n_max_boxes, align_metric) + if self._diagnostic_context is not None: + legal, small, boxes, anchors = self._diagnostic_context + counts = candidate_counts(result[2], overlaps, align_metric, legal, small, self.eps) + without_legal = small & ~legal.bool().any(-1) + counts.update(outside_capacity_counts(boxes, anchors, without_legal, result[1])) + self.records.append({"epoch": self.active_epoch, **counts}) + self._diagnostic_context = None + return result + + +def main(): + """Run one bounded diagnostic and save arguments, batch counts, and totals.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--data", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", default="cpu") + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--fraction", type=float, default=0.01) + parser.add_argument("--imgsz", type=int, default=800) + parser.add_argument("--batch", type=int, default=6) + args = parser.parse_args() + args.output = args.output.resolve() + if args.output.exists(): + raise FileExistsError(f"Refusing to overwrite diagnostic output: {args.output}") + if args.epochs < 1 or not 0 < args.fraction <= 1: + raise ValueError("epochs must be positive and fraction must be in (0, 1].") + + import ultralytics.utils.loss as loss_module + + original = loss_module.TaskAlignedAssigner + DiagnosticAssigner.records = [] + signatures = [] + yolo = YOLO(args.model) + + def start_epoch(trainer): + DiagnosticAssigner.active_epoch = int(trainer.epoch) + + def end_epoch(trainer): + DiagnosticAssigner.active_epoch = None + + def record_batch(trainer): + batch = trainer.batch + digest = hashlib.sha256() + for key in ("img", "bboxes", "cls", "batch_idx"): + digest.update(batch[key].detach().cpu().contiguous().numpy().tobytes()) + signatures.append({"files": batch["im_file"], "sha256": digest.hexdigest()}) + + yolo.add_callback("on_train_epoch_start", start_epoch) + yolo.add_callback("on_train_epoch_end", end_epoch) + yolo.add_callback("on_train_batch_start", record_batch) + try: + loss_module.TaskAlignedAssigner = DiagnosticAssigner + yolo.train( + data=args.data, + epochs=args.epochs, + fraction=args.fraction, + imgsz=args.imgsz, + batch=args.batch, + device=args.device, + workers=0, + pretrained=False, + optimizer="MuSGD", + amp=False, + seed=0, + deterministic=True, + mosaic=0, + close_mosaic=0, + val=False, + save=False, + plots=False, + patience=0, + project=str(args.output.parent), + name=args.output.name, + stal_candidate_mode="pure", + stal_enabled=False, + stal_zero_positive_rescue=False, + stal_nwd_weight=0, + ) + finally: + loss_module.TaskAlignedAssigner = original + DiagnosticAssigner.active_epoch = None + totals = Counter() + for record in DiagnosticAssigner.records: + totals.update({key: value for key, value in record.items() if key != "epoch"}) + if not DiagnosticAssigner.records: + raise RuntimeError("No training assignments observed; diagnostic is invalid.") + output = { + "protocol": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()}, + "torch": torch.__version__, + "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "totals": dict(totals), + "batches": DiagnosticAssigner.records, + "batch_signatures": signatures, + "interpretation": "Training-only candidate counts; no claim of AP improvement or official VisDrone evaluation.", + } + (args.output / "candidate-diagnostics.json").write_text(json.dumps(output, indent=2) + "\n", encoding="utf-8") + print(json.dumps({"totals": dict(totals)}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/stal/evaluate_scale_ap.py b/scripts/stal/evaluate_scale_ap.py new file mode 100644 index 000000000..d42ea2715 --- /dev/null +++ b/scripts/stal/evaluate_scale_ap.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python3 +"""Build COCO ground truth from YOLO labels and report size-binned AP for a detection split.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +from ultralytics.data.utils import IMG_FORMATS, check_det_dataset, img2label_paths + + +def resolve_split_images(split_source: str | list[str]) -> list[Path]: + """Resolve a dataset split expressed as image directories, files, or image-list text files.""" + sources = [split_source] if isinstance(split_source, (str, Path)) else split_source + images: list[Path] = [] + for source in sources: + path = Path(source) + if path.is_dir(): + images.extend(p for p in path.rglob("*") if p.is_file() and p.suffix[1:].lower() in IMG_FORMATS) + elif path.suffix.lower() == ".txt": + for line in path.read_text(encoding="utf-8").splitlines(): + item = line.strip() + if item: + images.append((path.parent / item).resolve() if item.startswith("./") else Path(item)) + elif path.is_file() and path.suffix[1:].lower() in IMG_FORMATS: + images.append(path) + else: + raise FileNotFoundError(f"Unable to resolve dataset split source: {path}") + return sorted(images) + + +def build_coco_ground_truth(image_files: list[Path], names: dict[int, str] | list[str]) -> dict[str, Any]: + """Convert YOLO normalized detection labels into an in-memory COCO ground-truth dictionary.""" + if not image_files: + raise ValueError("Evaluation split contains no images") + names = dict(enumerate(names)) if isinstance(names, list) else {int(k): v for k, v in names.items()} + label_files = img2label_paths([str(path) for path in image_files]) + images, annotations = [], [] + annotation_id = 1 + for image_path, label_path in zip(image_files, label_files): + with Image.open(image_path) as image: + width, height = image.size + stem = image_path.stem + image_id: int | str = int(stem) if stem.isnumeric() else stem + images.append({"id": image_id, "file_name": image_path.name, "width": width, "height": height}) + label_path = Path(label_path) + if not label_path.exists(): + raise FileNotFoundError(f"Missing evaluation label: {label_path}; use an empty file for a background image") + for line_number, line in enumerate(label_path.read_text(encoding="utf-8").splitlines(), start=1): + fields = line.split() + if not fields: + continue + if len(fields) != 5: + raise ValueError(f"Expected 5 YOLO fields in {label_path}:{line_number}, got {len(fields)}") + class_id, x_center, y_center, box_width, box_height = map(float, fields) + if not np.isfinite([class_id, x_center, y_center, box_width, box_height]).all(): + raise ValueError(f"Non-finite YOLO label in {label_path}:{line_number}") + if not class_id.is_integer() or not (0 <= x_center <= 1 and 0 <= y_center <= 1): + raise ValueError(f"Invalid class or normalized center in {label_path}:{line_number}") + if not (0 < box_width <= 1 and 0 < box_height <= 1): + raise ValueError(f"Invalid normalized box size in {label_path}:{line_number}") + class_id = int(class_id) + if class_id not in names: + raise ValueError(f"Class {class_id} in {label_path}:{line_number} is absent from dataset names") + box_width *= width + box_height *= height + x = x_center * width - box_width / 2 + y = y_center * height - box_height / 2 + annotations.append( + { + "id": annotation_id, + "image_id": image_id, + "category_id": class_id + 1, + "bbox": [x, y, box_width, box_height], + "area": box_width * box_height, + "iscrowd": 0, + } + ) + annotation_id += 1 + categories = [{"id": class_id + 1, "name": name} for class_id, name in sorted(names.items())] + return { + "info": {"description": "YOLO labels converted for STAL scale-aware evaluation"}, + "licenses": [], + "images": images, + "annotations": annotations, + "categories": categories, + } + + +def _mean_valid(values: np.ndarray) -> float: + """Average valid COCO precision/recall entries, preserving -1 for unavailable slices.""" + valid = values[values > -1] + return float(valid.mean()) if valid.size else -1.0 + + +def _coco_metric(evaluator, *, area: str = "all", iou: float | None = None, recall: bool = False) -> float: + """Read a COCO metric slice at maxDets=500 from accumulated precision or recall tensors.""" + params = evaluator.params + area_index = params.areaRngLbl.index(area) + max_det_index = params.maxDets.index(500) + if recall: + return _mean_valid(evaluator.eval["recall"][:, :, area_index, max_det_index]) + precision = evaluator.eval["precision"][:, :, :, area_index, max_det_index] + if iou is not None: + matches = np.flatnonzero(np.isclose(params.iouThrs, iou)) + if not matches.size: + raise ValueError(f"IoU threshold {iou} is absent from evaluator parameters") + precision = precision[matches] + return _mean_valid(precision) + + +def evaluate_scale_ap(annotation_path: Path, prediction_path: Path) -> dict[str, float]: + """Evaluate the project's COCO-style size metrics with the VisDrone maxDets=500 convention.""" + try: + from faster_coco_eval import COCO, COCOeval_faster + except ImportError as exc: + raise ImportError("Install faster-coco-eval>=1.6.7 before running scale-aware evaluation") from exc + + ground_truth = json.loads(annotation_path.read_text(encoding="utf-8")) + predictions = json.loads(prediction_path.read_text(encoding="utf-8")) + category_ids = {category["id"] for category in ground_truth["categories"]} + if len(category_ids) != len(ground_truth["categories"]): + raise ValueError("Duplicate category IDs in ground truth") + annotation_ids = set() + for record in ground_truth["annotations"]: + if record["id"] in annotation_ids: + raise ValueError(f"Duplicate annotation ID: {record['id']}") + annotation_ids.add(record["id"]) + if record["category_id"] not in category_ids: + raise ValueError(f"Ground truth refers to unknown category ID: {record['category_id']}") + bbox = record["bbox"] + if len(bbox) != 4 or not np.isfinite([*bbox, record["area"]]).all() or min(bbox[2:]) <= 0: + raise ValueError("Ground truth must have a finite area and a finite positive-size bbox") + if not np.isclose(record["area"], bbox[2] * bbox[3], rtol=1e-6, atol=1e-8): + raise ValueError("Project size evaluation requires GT area equal to bbox width * height") + for record in predictions: + if record["category_id"] not in category_ids: + raise ValueError(f"Prediction refers to unknown category ID: {record['category_id']}") + bbox = record["bbox"] + score = record["score"] + if len(bbox) != 4 or not np.isfinite([*bbox, score]).all() or min(bbox[2:]) <= 0 or not 0 <= score <= 1: + raise ValueError("Prediction must have a score in [0, 1] and a finite positive-size bbox") + # VisDrone stems contain underscores; the C++ evaluator requires integer image IDs. + # Remap both sides together, preserving the source JSON files for traceability. + image_ids = {image["id"]: index for index, image in enumerate(ground_truth["images"], start=1)} + if len(image_ids) != len(ground_truth["images"]): + raise ValueError("Duplicate image IDs in ground truth") + for record in ground_truth["images"]: + record["id"] = image_ids[record["id"]] + for record in ground_truth["annotations"]: + if record["image_id"] not in image_ids: + raise ValueError(f"Ground truth refers to unknown image ID: {record['image_id']}") + record["image_id"] = image_ids[record["image_id"]] + for record in predictions: + if record["image_id"] not in image_ids: + raise ValueError(f"Prediction refers to unknown image ID: {record['image_id']}") + record["image_id"] = image_ids[record["image_id"]] + annotation_api = COCO() + annotation_api.dataset = ground_truth + annotation_api.createIndex() + if predictions: + prediction_api = annotation_api.loadRes(predictions) + else: + # loadRes implementations may index the first detection; an empty run is valid evaluation input. + prediction_api = COCO() + prediction_api.dataset = {**ground_truth, "annotations": []} + prediction_api.createIndex() + evaluator = COCOeval_faster(annotation_api, prediction_api, iouType="bbox") + evaluator.params.imgIds = [image["id"] for image in annotation_api.dataset["images"]] + evaluator.params.maxDets = [1, 10, 500] + # COCO uses inclusive upper bounds; enforce the project's disjoint half-open bins. + evaluator.params.areaRng = [ + [0, 1e10], + [0, np.nextafter(float(32**2), -np.inf)], + [32**2, np.nextafter(float(96**2), -np.inf)], + [96**2, 1e10], + ] + evaluator.evaluate() + evaluator.accumulate() + return { + "AP": _coco_metric(evaluator), + "AP50": _coco_metric(evaluator, iou=0.50), + "AP75": _coco_metric(evaluator, iou=0.75), + "APs": _coco_metric(evaluator, area="small"), + "APm": _coco_metric(evaluator, area="medium"), + "APl": _coco_metric(evaluator, area="large"), + "AP50s": _coco_metric(evaluator, area="small", iou=0.50), + "AR500": _coco_metric(evaluator, recall=True), + "ARs500": _coco_metric(evaluator, area="small", recall=True), + "ARm500": _coco_metric(evaluator, area="medium", recall=True), + "ARl500": _coco_metric(evaluator, area="large", recall=True), + } + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data", required=True, help="Dataset YAML path") + parser.add_argument("--predictions", type=Path, required=True, help="Ultralytics predictions.json path") + parser.add_argument("--split", default="val", help="Dataset split key, usually val") + parser.add_argument("--annotations-out", type=Path, required=True, help="Generated COCO annotation JSON path") + parser.add_argument("--metrics-out", type=Path, required=True, help="Scale AP metrics JSON path") + return parser.parse_args() + + +def main() -> None: + """Build annotations, evaluate predictions, and save reproducible metric artifacts.""" + args = parse_args() + dataset = check_det_dataset(args.data, autodownload=False, split=args.split) + if args.split not in dataset: + raise KeyError(f"Split '{args.split}' is not defined in {args.data}") + image_files = resolve_split_images(dataset[args.split]) + ground_truth = build_coco_ground_truth(image_files, dataset["names"]) + args.annotations_out.parent.mkdir(parents=True, exist_ok=True) + args.annotations_out.write_text(json.dumps(ground_truth, ensure_ascii=False), encoding="utf-8") + metrics = evaluate_scale_ap(args.annotations_out, args.predictions) + metrics.update( + { + "images": len(ground_truth["images"]), + "annotations": len(ground_truth["annotations"]), + "max_dets": 500, + "metric_units": "fraction [0,1]; -1 means no evaluable GT; 0.01 equals 1 absolute AP point", + "iou_thresholds": "0.50:0.05:0.95", + "official_visdrone_metrics": False, + "protocol": "Project COCO-style supplemental analysis; these area bins are not VisDrone official bins", + "area_definition": "Original-image GT bbox area: small < 32^2, medium [32^2, 96^2), large >= 96^2 pixels", + } + ) + args.metrics_out.parent.mkdir(parents=True, exist_ok=True) + args.metrics_out.write_text(json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8") + print(json.dumps(metrics, indent=2, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/scripts/stal/export_visdrone_results.py b/scripts/stal/export_visdrone_results.py new file mode 100644 index 000000000..b6ac5de71 --- /dev/null +++ b/scripts/stal/export_visdrone_results.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""Convert Ultralytics predictions.json into the official VisDrone DET submission layout.""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import defaultdict +from pathlib import Path + +from ultralytics.data.utils import IMG_FORMATS + + +def export_visdrone_results(prediction_path: Path, image_dir: Path, output_dir: Path) -> dict[str, int]: + """Write one VisDrone DET result TXT per image, including empty files for images without detections.""" + predictions = json.loads(prediction_path.read_text(encoding="utf-8")) + if not isinstance(predictions, list): + raise TypeError("predictions JSON must contain a list") + image_paths = sorted(p for p in image_dir.iterdir() if p.is_file() and p.suffix[1:].lower() in IMG_FORMATS) + if not image_paths: + raise FileNotFoundError(f"No supported images found in {image_dir}") + known_stems = {path.stem for path in image_paths} + if len(known_stems) != len(image_paths): + raise ValueError("Duplicate image stems would overwrite the same result TXT") + numeric_stems = {int(stem): stem for stem in known_stems if stem.isnumeric()} + if len(numeric_stems) != sum(stem.isnumeric() for stem in known_stems): + raise ValueError("Ambiguous numeric image IDs in the image directory") + detections: dict[str, list[str]] = defaultdict(list) + for index, prediction in enumerate(predictions): + file_name = prediction.get("file_name") + image_id = prediction.get("image_id", "") + id_stem = str(image_id) + if id_stem not in known_stems and id_stem.isnumeric(): + id_stem = numeric_stems.get(int(id_stem), id_stem) + stem = Path(file_name).stem if file_name else id_stem + if stem not in known_stems: + raise ValueError(f"Prediction {index} refers to unknown image stem '{stem}'") + if file_name and "image_id" in prediction and id_stem != stem: + raise ValueError(f"Prediction {index} has conflicting file_name and image_id") + bbox = prediction.get("bbox") + if not isinstance(bbox, list) or len(bbox) != 4: + raise ValueError(f"Prediction {index} must contain bbox=[left, top, width, height]") + raw_category = prediction["category_id"] + category = float(raw_category) + if not math.isfinite(category) or not category.is_integer() or not 1 <= category <= 10: + raise ValueError(f"Prediction {index} has non-evaluated VisDrone category_id={category}; expected 1..10") + category = int(category) + left, top, width, height = (float(value) for value in bbox) + score = float(prediction["score"]) + if ( + not all(math.isfinite(value) for value in (left, top, width, height, score)) + or min(width, height) <= 0 + or not 0 <= score <= 1 + ): + raise ValueError(f"Prediction {index} must have a score in [0, 1] and a finite positive-size bbox") + detections[stem].append(f"{left:.3f},{top:.3f},{width:.3f},{height:.3f},{score:.8f},{category},-1,-1") + + # A directory from another split must not silently turn into a mixed official submission. + expected_files = {f"{stem}.txt" for stem in known_stems} + unexpected = [p.name for p in output_dir.glob("*") if p.suffix.lower() == ".txt" and p.name not in expected_files] + if unexpected: + raise ValueError( + f"Unexpected result TXT files in output directory: {sorted(unexpected)}; use a clean directory" + ) + output_dir.mkdir(parents=True, exist_ok=True) + for image_path in image_paths: + content = "\n".join(detections[image_path.stem]) + (output_dir / f"{image_path.stem}.txt").write_text(f"{content}\n" if content else "", encoding="utf-8") + return { + "images": len(image_paths), + "detections": len(predictions), + "empty_images": sum(not detections.get(path.stem) for path in image_paths), + } + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--predictions", type=Path, required=True, help="Ultralytics predictions.json") + parser.add_argument("--images", type=Path, required=True, help="Original VisDrone split image directory") + parser.add_argument("--output", type=Path, required=True, help="Official-toolkit result TXT directory") + return parser.parse_args() + + +def main() -> None: + """Export predictions and print a compact manifest.""" + args = parse_args() + summary = export_visdrone_results(args.predictions, args.images, args.output) + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/test_stal_amp_diagnostic_contract.py b/tests/test_stal_amp_diagnostic_contract.py new file mode 100644 index 000000000..ed9016c01 --- /dev/null +++ b/tests/test_stal_amp_diagnostic_contract.py @@ -0,0 +1,51 @@ +"""Regression checks for shared training initialization and complete batch fingerprints.""" + +import pytest +import torch + +from scripts.stal.diagnose_amp_training import ( + batch_signature, + first_batch_mismatch, + shared_state_trainer, + state_digest, + tensor_digest, +) +from ultralytics.models.yolo.detect import DetectionTrainer + + +def test_shared_state_restores_new_training_model(monkeypatch): + """Dataset-specific rebuilt models must use the first arm's actual state, including buffers.""" + monkeypatch.setattr(DetectionTrainer, "get_model", lambda *a, **kw: torch.nn.BatchNorm1d(10)) + shared = {} + trainer_type = shared_state_trainer(shared) + first = trainer_type.__new__(trainer_type).get_model() + shared["weight"].fill_(7) + shared["running_mean"].fill_(3) + second = trainer_type.__new__(trainer_type).get_model() + assert torch.all(second.weight == 7) + assert torch.all(second.running_mean == 3) + assert state_digest(second.state_dict()) == state_digest(shared) + assert torch.all(first.weight == 1) # the captured state owns its storage + + +@pytest.mark.parametrize("key", ["img", "bboxes", "cls", "batch_idx"]) +def test_batch_signature_detects_equal_sum_permutation(key): + """Old aggregate sums miss pixel, target, class and image-association permutations.""" + batch = {name: torch.tensor([1, 2]) for name in ("img", "bboxes", "cls", "batch_idx")} + changed = {**batch, key: batch[key].flip(0)} + assert batch[key].sum() == changed[key].sum() + assert batch_signature(batch) != batch_signature(changed) + + +def test_tensor_digest_preserves_shape_dtype_and_empty_tensors(): + """Identical bytes with different interpretation cannot count as identical inputs.""" + value = torch.tensor([1, 2], dtype=torch.int32) + assert tensor_digest(value) != tensor_digest(value.reshape(1, 2)) + assert tensor_digest(value) != tensor_digest(value.view(torch.float32)) + assert tensor_digest(torch.empty(0)) == tensor_digest(torch.empty(0)) + + +@pytest.mark.parametrize("left,right,expected", [([], [], None), (["a"], [], 0), (["a"], ["a", "b"], 1)]) +def test_batch_mismatch_detects_length_difference(left, right, expected): + """A missing suffix must have a concrete first mismatch index.""" + assert first_batch_mismatch(left, right) == expected diff --git a/tests/test_stal_amp_precision.py b/tests/test_stal_amp_precision.py new file mode 100644 index 000000000..eaeba00b6 --- /dev/null +++ b/tests/test_stal_amp_precision.py @@ -0,0 +1,109 @@ +"""Exercise actual autocast head outputs through assignment and complete detection loss.""" + +from types import SimpleNamespace + +import pytest +import torch + +from ultralytics.cfg import get_cfg +from ultralytics.nn.tasks import DetectionModel # noqa: F401 - normal loss import initialization +from ultralytics.utils.loss import v8DetectionLoss +from ultralytics.utils.tal import TaskAlignedAssigner + + +@pytest.mark.parametrize("mode", ["pure", "fixed", "adaptive"]) +@pytest.mark.parametrize("small_scene", [False, True]) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA unavailable; CPU BF16 is not a substitute for CUDA FP16" + ), + ), + ], +) +def test_actual_autocast_assignment_and_detection_gradients(mode, device, small_scene): + """Compare masks, top-k, counts, three losses and head gradients on a deterministic non-tied scene.""" + torch.manual_seed(14) + head = torch.nn.Conv1d(4, 65, 1).to(device) + if small_scene: + # Decode boxes on the same pixel scale as the tiny GTs; random wide boxes can have zero CIoU. + with torch.no_grad(): + head.weight[:64].mul_(0.1) + bias = head.bias[:64].reshape(4, 16) + bias.fill_(-4.0) + bias[:, 0] = 4.0 + bias[:, 1] = 3.0 + head.args = get_cfg( + overrides={ + "stal_candidate_mode": mode, + "stal_warmup_epochs": 0.0, + **({"stal_small_topk": 6, "stal_small_topk_min": 6} if small_scene else {}), + } + ) + head.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0], device=device), nc=1, reg_max=16)] + features = torch.randn(1, 4, 84, device=device) + batch = { + "batch_idx": torch.tensor([0.0, 0.0], device=device), + "cls": torch.zeros(2, 1, device=device), + "bboxes": torch.tensor( + [[0.3125, 0.3125, 0.05, 0.05], [0.6875, 0.6875, 0.06, 0.06]] + if small_scene + else [[0.25, 0.25, 0.15, 0.15], [0.70, 0.70, 0.2, 0.2]], + device=device, + ), + "epoch": 10, + } + results = [] + for amp in (False, True): + head.zero_grad(set_to_none=True) + criterion = v8DetectionLoss(head) + selected = [] + original = criterion.assigner.get_pos_mask + + def capture(*args, original=original, selected=selected, criterion=criterion, **kwargs): + if small_scene: + boxes, anchors, mask = args[3:6] + gate = criterion.assigner.small_target_mask(boxes, kwargs["image_size"]) + assert gate.all(), "Fixture must activate the small-target area gate" + expanded = criterion.assigner.select_candidates_in_gts(anchors, boxes, mask, **kwargs) + base = criterion.assigner.select_candidates_in_gts( + anchors, boxes, mask, **kwargs, relaxation_override=0.0 + ) + if mode == "adaptive": + assert (expanded & ~base).any(), "Adaptive expansion must add candidates beyond fixed STAL" + result = original(*args, **kwargs) + selected.append(result[0].detach().clone()) + return result + + criterion.assigner.get_pos_mask = capture + dtype = torch.float16 if device == "cuda" else torch.bfloat16 + with torch.autocast(device, dtype=dtype, enabled=amp): + output = head(features) + assert output.dtype == (dtype if amp else torch.float32) + predictions = { + "boxes": output[:, :64], + "scores": output[:, 64:], + "feats": [torch.zeros(1, 1, s, s, device=device) for s in (8, 4, 2)], + } + assigned, loss, _ = criterion.get_assigned_targets_and_loss(predictions, batch) + loss.sum().backward() + gradients = torch.cat([p.grad.flatten() for p in head.parameters()]) + assert torch.isfinite(loss).all() and torch.isfinite(gradients).all() + assert assigned[0].any() + results.append((assigned[0].clone(), assigned[1].clone(), selected[0], loss.detach(), gradients.clone())) + for index in range(3): + assert torch.equal(results[0][index], results[1][index]) + # BF16 is coarser than FP16. This bounded tolerance is for the synthetic fixture, not a training-equivalence claim. + for index in (3, 4): + torch.testing.assert_close(results[0][index], results[1][index], rtol=0.05, atol=0.01) + + +def test_adaptive_amp_scene_rejects_disabled_expansion(monkeypatch): + """The new fixture must fail if the adaptive geometry silently becomes fixed STAL.""" + monkeypatch.setattr(TaskAlignedAssigner, "stal_relaxation_at_epoch", lambda self, epoch: 0.0) + with pytest.raises(AssertionError, match="Adaptive expansion must add candidates"): + test_actual_autocast_assignment_and_detection_gradients("adaptive", "cpu", True) diff --git a/tests/test_stal_area_precision.py b/tests/test_stal_area_precision.py new file mode 100644 index 000000000..1861ef252 --- /dev/null +++ b/tests/test_stal_area_precision.py @@ -0,0 +1,68 @@ +"""Guard scale gates and telemetry at the formal 800-pixel training resolution.""" + +from types import SimpleNamespace + +import torch + +from ultralytics.cfg import get_cfg +from ultralytics.nn.tasks import DetectionModel # noqa: F401 +from ultralytics.utils.loss import v8DetectionLoss +from ultralytics.utils.tal import TaskAlignedAssigner + + +def test_half_canvas_does_not_turn_medium_targets_into_stal_targets(): + """800*800 overflows FP16: 100*100/inf must not become a false small-target classification.""" + assigner = TaskAlignedAssigner() + boxes = torch.tensor([[[0.0, 0.0, 20.0, 20.0], [0.0, 0.0, 80.0, 80.0], [0.0, 0.0, 100.0, 100.0]]]) + size = torch.tensor([800.0, 800.0], dtype=torch.float16) + assert assigner.small_target_mask(boxes, size).tolist() == [[True, False, False]] + + +def test_half_candidate_geometry_matches_fp32_at_formal_resolution(): + """A 100px box is above 1% at 800px and must not receive adaptive expansion.""" + assigner = TaskAlignedAssigner(stal_candidate_mode="adaptive", stal_warmup_epochs=0) + boxes = torch.tensor([[[100.0, 100.0, 200.0, 200.0]]]) + anchors = torch.tensor([[98.0, 150.0], [150.0, 150.0]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + for dtype in (torch.float32, torch.float16): + selected = assigner.select_candidates_in_gts(anchors.to(dtype), boxes.to(dtype), valid, image_size=(800, 800)) + assert selected.tolist() == [[[False, True]]] + + +def test_half_assignment_statistics_preserve_stal_area_gate(): + """A medium-sized target must not enter the relative-area STAL statistics group.""" + assigner = TaskAlignedAssigner() + stats = assigner.assignment_statistics( + torch.tensor([[[100.0, 100.0, 200.0, 200.0]]], dtype=torch.float16), + torch.ones(1, 1, 1, dtype=torch.bool), + torch.ones(1, 1, dtype=torch.bool), + torch.zeros(1, 1, dtype=torch.long), + (800, 800), + ) + assert stats[:3].tolist() == [0, 0, 0] + assert stats[-3:].tolist() == [1, 1, 0] + + +def test_fp16_loss_score_statistics_exclude_non_stal_gt(): + """Exercise the actual loss telemetry with an FP16 800px canvas and FP32 training labels.""" + model = torch.nn.Linear(1, 1) + model.args = get_cfg(overrides={"stal_stats": True}) + model.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=1, reg_max=16)] + loss = v8DetectionLoss(model) + predictions = { + "feats": [torch.zeros(1, 1, s, s) for s in (100, 50, 25)], + "boxes": torch.zeros(1, 64, 13125, dtype=torch.float16), + "scores": torch.zeros(1, 1, 13125, dtype=torch.float16), + } + batch = { + "batch_idx": torch.zeros(1), + "cls": torch.zeros(1, 1), + "bboxes": torch.tensor([[0.5, 0.5, 0.125, 0.125]]), + "epoch": 10, + } + # CPU autocast keeps the BCE computation supported while predictions remain genuinely FP16. + with torch.autocast("cpu", dtype=torch.bfloat16): + loss.get_assigned_targets_and_loss(predictions, batch) + stats = loss.pop_target_score_stats() + assert stats[3] > 0 + assert stats[:3].tolist() == [0.0, 0.0, 0.0] diff --git a/tests/test_stal_assigner.py b/tests/test_stal_assigner.py new file mode 100644 index 000000000..407970b9e --- /dev/null +++ b/tests/test_stal_assigner.py @@ -0,0 +1,1190 @@ +# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license + +import pytest +import torch + +from ultralytics.cfg import check_cfg, get_cfg +from ultralytics.utils.tal import TaskAlignedAssigner + + +def _candidate_fixture(): + """Return anchors and a 10x8 box that does not trigger the locked baseline's <8 pixel rule.""" + anchors = torch.tensor([[4.0, 4.0], [12.0, 4.0], [20.0, 4.0]]) + boxes = torch.tensor([[[7.0, 0.0, 17.0, 8.0]]]) + valid = torch.ones((1, 1, 1)) + return anchors, boxes, valid + + +def test_stal_disabled_preserves_locked_candidate_selection(): + """Explicitly disabling STAL must match the locked baseline's default behavior.""" + anchors, boxes, valid = _candidate_fixture() + baseline = TaskAlignedAssigner(stride=[8, 16, 32]) + disabled = TaskAlignedAssigner(stride=[8, 16, 32], stal_enabled=False) + + expected = baseline.select_candidates_in_gts(anchors, boxes, valid) + actual = disabled.select_candidates_in_gts(anchors, boxes, valid, image_size=(100, 100), epoch=10) + + assert torch.equal(actual, expected) + assert actual.tolist() == [[[False, True, False]]] + + +def test_pure_tal_uses_unmodified_ground_truth_region(): + """Pure TAL must not inherit the repository's fixed-stride candidate expansion.""" + anchors = torch.tensor([[4.0, 4.0], [12.0, 4.0]]) + boxes = torch.tensor([[[7.0, 1.0, 9.0, 7.0]]]) + valid = torch.ones((1, 1, 1)) + + candidates = TaskAlignedAssigner(stal_candidate_mode="pure").select_candidates_in_gts(anchors, boxes, valid) + + assert candidates.tolist() == [[[False, False]]] + + +def test_fixed_mode_preserves_locked_stride_expansion(): + """The default fixed mode must retain the locked repository behavior for sub-stride boxes.""" + anchors = torch.tensor([[4.0, 4.0], [12.0, 4.0]]) + boxes = torch.tensor([[[7.0, 1.0, 9.0, 7.0]]]) + valid = torch.ones((1, 1, 1)) + + candidates = TaskAlignedAssigner(stal_candidate_mode="fixed").select_candidates_in_gts(anchors, boxes, valid) + + assert candidates.tolist() == [[[True, True]]] + + +def test_stal_area_rule_expands_candidate_region_after_warmup(): + """A relative-area target should admit extra anchors once symmetric relaxation is active.""" + anchors, boxes, valid = _candidate_fixture() + assigner = TaskAlignedAssigner( + stride=[8, 16, 32], + stal_enabled=True, + stal_area_threshold=0.01, + stal_relaxation=8.0, + stal_warmup_epochs=10, + ) + + candidates = assigner.select_candidates_in_gts(anchors, boxes, valid, image_size=(100, 100), epoch=10) + + assert candidates.tolist() == [[[True, True, True]]] + + +def test_stal_trigger_uses_original_area_before_locked_width_expansion(): + """The locked ciou_quality[0, 0, 0] + assert mixed_quality[0, 1, 1] == pytest.approx(ciou_quality[0, 1, 1]) + assert torch.equal(ciou_quality[~legal], mixed_quality[~legal]) + assert torch.equal(mixed_quality, mixed_targets) + + +def test_nwd_ciou_target_mode_decouples_matching_from_supervision_quality(): + """NWD may rank tiny candidates while CIoU continues to set the final target-score scale.""" + scores = torch.full((1, 1, 1), 0.8) + predicted = torch.tensor([[[11.0, 10.0, 13.0, 12.0]]]) + labels = torch.zeros((1, 1, 1)) + boxes = torch.tensor([[[10.0, 10.0, 12.0, 12.0]]]) + legal = torch.ones((1, 1, 1), dtype=torch.bool) + assigner = TaskAlignedAssigner(num_classes=1, stal_nwd_weight=0.25, stal_nwd_target_mode="ciou") + assigner.bs = assigner.n_max_boxes = 1 + + _, match_quality, target_quality = assigner.get_box_metrics( + scores, predicted, labels, boxes, legal, image_size=(100, 100) + ) + + assert match_quality.item() > target_quality.item() + assert target_quality.item() == pytest.approx(assigner.iou_calculation(boxes[0], predicted[0]).item()) + + +def test_nwd_ciou_target_mode_can_select_zero_supervision_candidate(): + """A zero-CIoU candidate may be matched by NWD but receive no target-score supervision in ciou mode.""" + scores = torch.full((1, 1, 1), 0.8) + predicted = torch.tensor([[[12.0, 10.0, 14.0, 12.0]]]) + labels = torch.zeros((1, 1, 1)) + boxes = torch.tensor([[[10.0, 10.0, 12.0, 12.0]]]) + anchors = torch.tensor([[12.0, 11.0]]) + valid = torch.ones((1, 1, 1), dtype=torch.bool) + assigner = TaskAlignedAssigner(num_classes=1, topk=1, stal_nwd_weight=0.25, stal_nwd_target_mode="ciou") + + _, _, target_scores, foreground, _ = assigner( + scores, predicted, anchors, labels, boxes, valid, image_size=(100, 100) + ) + + assert foreground.item() + assert target_scores.sum().item() == 0.0 + + +def test_nwd_weighted_target_mode_restores_bounded_supervision(): + """A smaller target blend should supervise NWD matches without using the full matching quality.""" + scores = torch.full((1, 1, 1), 0.8) + predicted = torch.tensor([[[12.0, 10.0, 14.0, 12.0]]]) + labels = torch.zeros((1, 1, 1)) + boxes = torch.tensor([[[10.0, 10.0, 12.0, 12.0]]]) + legal = torch.ones((1, 1, 1), dtype=torch.bool) + assigner = TaskAlignedAssigner( + num_classes=1, + stal_nwd_weight=0.25, + stal_nwd_target_mode="weighted", + stal_nwd_target_weight=0.05, + ) + assigner.bs = assigner.n_max_boxes = 1 + + _, match_quality, target_quality = assigner.get_box_metrics( + scores, predicted, labels, boxes, legal, image_size=(100, 100) + ) + + assert 0.0 < target_quality.item() < match_quality.item() + + +def test_nwd_zero_score_floor_only_changes_selected_zero_score_small_anchor(): + """The NWD floor should activate after matching without replacing positive CIoU target scores.""" + scores = torch.full((1, 2, 1), 0.8) + predicted = torch.tensor([[[12.0, 10.0, 14.0, 12.0], [20.5, 20.0, 22.5, 22.0]]]) + labels = torch.zeros((1, 2, 1)) + boxes = torch.tensor([[[10.0, 10.0, 12.0, 12.0], [20.0, 20.0, 22.0, 22.0]]]) + anchors = torch.tensor([[12.0, 11.0], [21.5, 21.0]]) + valid = torch.ones((1, 2, 1), dtype=torch.bool) + assigner = TaskAlignedAssigner( + num_classes=1, + topk=1, + stal_nwd_weight=0.25, + stal_nwd_target_mode="ciou", + stal_nwd_zero_score_floor=0.01, + ) + + _, _, target_scores, foreground, target_gt_idx = assigner( + scores, predicted, anchors, labels, boxes, valid, image_size=(100, 100) + ) + + assigned_scores = target_scores.sum(-1)[foreground] + assigned_gt = target_gt_idx[foreground] + assert assigned_scores[assigned_gt == 0].item() == pytest.approx(0.01) + assert assigned_scores[assigned_gt == 1].item() > 0.01 + + +@pytest.mark.parametrize( + "config", + [ + {"stal_enabled": "true"}, + {"stal_stats": "true"}, + {"stal_candidate_mode": 1}, + {"stal_area_threshold": "0.01"}, + {"stal_small_topk": "13"}, + {"stal_small_topk_min": "3"}, + {"stal_min_base_candidates": "1"}, + {"stal_max_extra_candidates": "1"}, + {"stal_min_candidate_guarantee": "true"}, + {"stal_expanded_quality_ratio": "0.5"}, + {"stal_expanded_score_floor": "0.5"}, + {"stal_relaxation": "8"}, + {"stal_relaxation_scale_mode": 1}, + {"stal_relaxation_min": "4"}, + {"stal_min_size_stride_ratio": "2"}, + {"stal_warmup_epochs": "10"}, + {"stal_zero_positive_rescue": "true"}, + {"stal_rescue_score_floor": "0.05"}, + {"stal_rescue_floor_decay_epochs": "10"}, + ], +) +def test_stal_cli_configuration_is_type_checked(config): + """String-like programmatic values must not silently bypass configuration validation.""" + with pytest.raises(TypeError): + check_cfg(config) + + +@pytest.mark.parametrize( + "config", + [ + {"stal_candidate_mode": "unknown"}, + {"stal_area_threshold": 0.0}, + {"stal_small_topk": 0}, + {"stal_small_topk_min": 0}, + {"stal_small_topk": 4, "stal_small_topk_min": 5}, + {"stal_min_base_candidates": -1}, + {"stal_max_extra_candidates": -1}, + {"stal_expanded_quality_ratio": -0.1}, + {"stal_expanded_quality_ratio": 1.1}, + {"stal_expanded_score_floor": -0.1}, + {"stal_expanded_score_floor": 1.1}, + {"stal_relaxation": -1.0}, + {"stal_relaxation_scale_mode": "unknown"}, + {"stal_relaxation_min": -1.0}, + {"stal_relaxation": 4.0, "stal_relaxation_min": 5.0}, + {"stal_min_size_stride_ratio": -1.0}, + {"stal_min_size_stride_ratio": 2.0}, + { + "stal_relaxation": 0.0, + "stal_relaxation_scale_mode": "sqrt_area", + "stal_min_size_stride_ratio": 2.0, + }, + {"stal_warmup_epochs": -1.0}, + {"stal_enabled": True, "stal_candidate_mode": "pure"}, + {"stal_rescue_score_floor": 0.05}, + { + "stal_candidate_mode": "adaptive", + "stal_zero_positive_rescue": True, + "stal_rescue_floor_decay_epochs": -1.0, + }, + ], +) +def test_stal_central_contract_rejects_invalid_values_and_relationships(config): + """CLI configuration validation must fail before constructing a model or starting training.""" + with pytest.raises(ValueError): + check_cfg(config) diff --git a/tests/test_stal_candidate_capacity.py b/tests/test_stal_candidate_capacity.py new file mode 100644 index 000000000..770fdfa67 --- /dev/null +++ b/tests/test_stal_candidate_capacity.py @@ -0,0 +1,89 @@ +"""Verify the opt-in geometric capacity filter for expanded small-target candidates.""" + +import pytest +import torch + +from ultralytics.cfg import check_cfg, get_cfg +from ultralytics.utils.tal import TaskAlignedAssigner + + +def test_capacity_keeps_inside_and_feasible_outside_candidates(): + """Keep the 0.75 boundary but reject the 0.60-capacity point that fixed STAL admitted.""" + points = torch.tensor([[3.0, 3.0], [8.0, 3.0], [10.0, 3.0]]) + boxes = torch.tensor([[[0.0, 0.0, 6.0, 6.0]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + before = boxes.clone() + fixed = TaskAlignedAssigner(stal_candidate_mode="fixed") + capped = TaskAlignedAssigner(stal_candidate_mode="fixed", stal_candidate_iou_floor=0.75) + assert fixed.select_candidates_in_gts(points, boxes, valid).tolist() == [[[True, True, True]]] + assert capped.select_candidates_in_gts(points, boxes, valid, image_size=(100, 100)).tolist() == [ + [[True, True, False]] + ] + assert torch.equal(boxes, before) + + +def test_capacity_preserves_non_small_and_pure_candidates(): + """Non-small GTs are unaffected; pure TAL ignores the expansion-only filter.""" + points = torch.tensor([[3.0, 3.0], [8.0, 3.0], [10.0, 3.0]]) + boxes = torch.tensor([[[0.0, 0.0, 6.0, 200.0]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + for mode in ("pure", "fixed", "adaptive"): + original = TaskAlignedAssigner(stal_candidate_mode=mode) + filtered = TaskAlignedAssigner(stal_candidate_mode=mode, stal_candidate_iou_floor=0.75) + a = original.select_candidates_in_gts(points, boxes, valid, image_size=(100, 100)) + b = filtered.select_candidates_in_gts(points, boxes, valid, image_size=(100, 100)) + assert torch.equal(a, b) + + +@pytest.mark.parametrize("value", [-0.1, 1.1, float("nan")]) +def test_capacity_rejects_invalid_floor(value): + """Reject invalid thresholds both at configuration and assigner boundaries.""" + with pytest.raises(ValueError): + TaskAlignedAssigner(stal_candidate_iou_floor=value) + with pytest.raises(ValueError): + check_cfg({"stal_candidate_iou_floor": value}) + + +def test_capacity_default_and_type_contract(): + """Default is disabled and string-like programmatic thresholds are rejected.""" + assert get_cfg().stal_candidate_iou_floor == 0.0 + with pytest.raises(TypeError): + check_cfg({"stal_candidate_iou_floor": "0.75"}) + + +def test_capacity_requires_image_size_when_enabled(): + """A missing scale must fail rather than silently apply the small-target gate incorrectly.""" + capped = TaskAlignedAssigner(stal_candidate_mode="fixed", stal_candidate_iou_floor=0.75) + with pytest.raises(ValueError, match="image_size"): + capped.select_candidates_in_gts( + torch.tensor([[3.0, 3.0]]), torch.tensor([[[0.0, 0.0, 6.0, 6.0]]]), torch.ones(1, 1, 1) + ) + + +@pytest.mark.parametrize( + "bbox,point,relaxation,floor", + [ + ([100.0, 100.0, 120.0, 120.0], [500.0, 500.0], 800.0, 0.002), + ([0.0, 0.0, 0.0001, 0.0001], [0.00005, 0.00005], 0.0, 0.75), + ], +) +def test_capacity_half_precision_matches_fp32(bbox, point, relaxation, floor): + """Neither overflowing enclosing areas nor underflowing tiny GT areas may reject feasible candidates.""" + assigner = TaskAlignedAssigner( + stal_candidate_mode="adaptive", + stal_relaxation=relaxation, + stal_warmup_epochs=0, + stal_candidate_iou_floor=floor, + ) + masks = [] + for dtype in (torch.float32, torch.float16): + masks.append( + assigner.select_candidates_in_gts( + torch.tensor([point], dtype=dtype), + torch.tensor([[bbox]], dtype=dtype), + torch.ones(1, 1, 1, dtype=torch.bool), + image_size=(800, 800), + ) + ) + assert masks[0].item() + assert torch.equal(masks[0], masks[1]) diff --git a/tests/test_stal_candidate_diagnostics.py b/tests/test_stal_candidate_diagnostics.py new file mode 100644 index 000000000..2008ec91d --- /dev/null +++ b/tests/test_stal_candidate_diagnostics.py @@ -0,0 +1,56 @@ +"""Protect candidate diagnostic semantics and unchanged assignment outputs.""" + +import torch + +from scripts.stal.diagnose_rescue_candidates import DiagnosticAssigner, candidate_counts, outside_capacity_counts +from ultralytics.utils.tal import TaskAlignedAssigner + + +def test_outside_anchor_capacity_and_occupied_candidate(): + """A 2x2 GT at the origin cannot reach AP50 with its only point at (4,4).""" + boxes = torch.tensor([[[0.0, 0.0, 2.0, 2.0], [3.0, 3.0, 5.0, 5.0]]]) + anchors = torch.tensor([[4.0, 4.0], [12.0, 12.0]]) + selected = torch.ones(1, 2, dtype=torch.bool) + counts = outside_capacity_counts(boxes, anchors, selected, torch.tensor([[True, False]])) + assert counts["outside_gt"] == 2 + assert counts["outside_any_iou_cap_ge_0.5"] == 1 + assert counts["outside_any_iou_cap_ge_0.95"] == 1 + assert counts["outside_free_iou_cap_ge_0.5"] == 0 + + +def test_counts_separate_no_candidate_occupied_and_threshold_suppression(): + """A positive CIoU can be filtered by alignment eps; occupied anchors cannot be rescued.""" + mask = torch.tensor([[[1, 0], [0, 0], [0, 0], [0, 0]]]) + legal = torch.tensor([[[1, 0], [0, 1], [1, 0], [0, 0]]], dtype=torch.bool) + quality = torch.tensor([[[1.0, 0.0], [0.0, 0.03], [0.8, 0.0], [0.0, 0.0]]]) + metric = quality.pow(6) * 0.1 + counts = candidate_counts(mask, quality, metric, legal, torch.ones(1, 4, dtype=torch.bool), 1e-9) + assert counts["uncovered"] == 3 + assert counts["uncovered_without_legal"] == 1 + assert counts["uncovered_with_free_legal"] == 1 + assert counts["positive_ciou_but_alignment_below_eps"] == 1 + assert counts["suppressed_ciou_ge_0.03"] == 1 + assert counts["suppressed_ciou_ge_0.05"] == 0 + + +def test_diagnostic_preserves_all_assigner_outputs(): + """Observation must not change pure TAL labels, masks, boxes, scores, or GT ownership.""" + inputs = ( + torch.tensor([[[0.5], [0.1]]]), + torch.tensor([[[0.0, 0.0, 6.0, 6.0], [0.0, 0.0, 10.0, 6.0]]]), + torch.tensor([[3.0, 3.0], [8.0, 3.0]]), + torch.zeros(1, 2, 1), + torch.tensor([[[0.0, 0.0, 6.0, 6.0], [0.0, 0.0, 10.0, 6.0]]]), + torch.ones(1, 2, 1, dtype=torch.bool), + ) + baseline = TaskAlignedAssigner(topk=1, num_classes=1, stal_candidate_mode="pure") + diagnostic = DiagnosticAssigner(topk=1, num_classes=1, stal_candidate_mode="pure") + DiagnosticAssigner.active_epoch = 0 + DiagnosticAssigner.records = [] + try: + expected = baseline(*inputs, image_size=(100, 100)) + actual = diagnostic(*inputs, image_size=(100, 100)) + assert all(torch.equal(a, b) for a, b in zip(expected, actual)) + assert len(DiagnosticAssigner.records) == 1 + finally: + DiagnosticAssigner.active_epoch = None diff --git a/tests/test_stal_empty_batch_loss.py b/tests/test_stal_empty_batch_loss.py new file mode 100644 index 000000000..4b3a8d537 --- /dev/null +++ b/tests/test_stal_empty_batch_loss.py @@ -0,0 +1,83 @@ +"""Exercise background-only assignment through detection loss and diagnostics.""" + +from types import SimpleNamespace + +import pytest +import torch + +from ultralytics.cfg import get_cfg +from ultralytics.nn.tasks import DetectionModel # noqa: F401 - initialize the normal model/loss import path +from ultralytics.utils.loss import v8DetectionLoss +from ultralytics.utils.tal import TaskAlignedAssigner + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_empty_assignment_returns_indexable_types(dtype): + """An empty target batch must preserve discrete output types and prediction precision.""" + scores = torch.zeros(2, 20, 3, dtype=dtype) + boxes = torch.zeros(2, 20, 4, dtype=dtype) + labels, targets, weights, foreground, indices = TaskAlignedAssigner(num_classes=3)( + scores, + boxes, + torch.zeros(20, 2), + torch.empty(2, 0, 1), + torch.empty(2, 0, 4), + torch.empty(2, 0, 1, dtype=torch.bool), + ) + assert labels.dtype == indices.dtype == torch.long + assert foreground.dtype == torch.bool + assert targets.dtype == weights.dtype == dtype + assert labels.shape == indices.shape == foreground.shape == (2, 20) + assert (labels == 3).all() + assert not foreground.any() + assert not indices.any() + assert not targets.any() + assert not weights.any() + + +@pytest.mark.parametrize("mode", ["pure", "fixed", "adaptive", "rescue", "quality"]) +@pytest.mark.parametrize("stats", [False, True]) +def test_empty_batch_detection_loss_and_backward(mode, stats): + """Diagnostics must not break background classification supervision or accumulate positive counts.""" + config = { + "stal_candidate_mode": mode if mode in {"pure", "fixed"} else "adaptive", + "stal_stats": stats, + "stal_zero_positive_rescue": mode == "rescue", + "stal_rescue_score_floor": 0.01 if mode == "rescue" else 0.0, + "stal_expanded_quality_ratio": 0.25 if mode == "quality" else 0.0, + } + model = torch.nn.Linear(1, 1) + model.args = get_cfg(overrides=config) + model.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=3, reg_max=16)] + criterion = v8DetectionLoss(model) + predictions = { + "feats": [torch.zeros(2, 1, size, size) for size in (8, 4, 2)], + "boxes": torch.zeros(2, 64, 84, requires_grad=True), + "scores": torch.zeros(2, 3, 84, requires_grad=True), + } + batch = { + "batch_idx": torch.empty(0), + "cls": torch.empty(0, 1), + "bboxes": torch.empty(0, 4), + "epoch": 10, + } + assigned, loss, _ = criterion.get_assigned_targets_and_loss(predictions, batch) + assert assigned[0].dtype == torch.bool + assert assigned[1].dtype == torch.long + assert not assigned[0].any() + expected_cls = ( + torch.nn.functional.binary_cross_entropy_with_logits( + predictions["scores"], torch.zeros_like(predictions["scores"]), reduction="sum" + ) + * model.args.cls + ) + torch.testing.assert_close(loss[1], expected_cls) + assert loss[0] == loss[2] == 0 + loss.sum().backward() + assert torch.isfinite(predictions["scores"].grad).all() + assert (predictions["scores"].grad > 0).all() + assert predictions["boxes"].grad is None # No regression targets in a background-only batch. + assert not criterion.pop_assignment_stats().any() + assert not criterion.pop_rescue_stats().any() + assert not criterion.pop_assignment_stage_stats().any() + assert not criterion.pop_target_score_stats().any() diff --git a/tests/test_stal_numeric_config.py b/tests/test_stal_numeric_config.py new file mode 100644 index 000000000..a4d851693 --- /dev/null +++ b/tests/test_stal_numeric_config.py @@ -0,0 +1,56 @@ +"""Check numerical experiment parameters cannot silently turn mechanisms into no-ops.""" + +import pytest + +from ultralytics.cfg import CFG_FLOAT_KEYS, CFG_FRACTION_KEYS, CFG_INT_KEYS, get_cfg +from ultralytics.utils.tal import TaskAlignedAssigner + + +@pytest.mark.parametrize("key", sorted(k for k in CFG_FLOAT_KEYS | CFG_FRACTION_KEYS if k.startswith("stal_"))) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf"), True]) +def test_stal_numeric_parameters_reject_nonfinite_and_boolean(key, value): + """NaN comparisons and bool-as-int must not bypass the central configuration contract.""" + with pytest.raises((ValueError, TypeError)): + get_cfg(overrides={key: value}) + + +@pytest.mark.parametrize("key", sorted(k for k in CFG_INT_KEYS if k.startswith("stal_"))) +def test_stal_integer_parameters_reject_boolean(key): + """Python booleans are integers but are not experiment counts.""" + with pytest.raises(TypeError): + get_cfg(overrides={key: True}) + + +@pytest.mark.parametrize( + "key", + [ + "stal_relaxation", + "stal_min_size_stride_ratio", + "stal_warmup_epochs", + "stal_rescue_floor_decay_epochs", + "stal_nwd_constant", + ], +) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_direct_assigner_rejects_nonfinite_unbounded_parameters(key, value): + """Research scripts constructing the assigner directly must also reject non-finite parameters.""" + with pytest.raises(ValueError, match="finite"): + TaskAlignedAssigner(**{key: value}) + + +@pytest.mark.parametrize( + "key", sorted(k for k in CFG_FLOAT_KEYS | CFG_FRACTION_KEYS | CFG_INT_KEYS if k.startswith("stal_")) +) +def test_direct_assigner_rejects_boolean_numeric_parameters(key): + """Direct construction must preserve the same bool-versus-number boundary as central configuration.""" + with pytest.raises((TypeError, ValueError)): + TaskAlignedAssigner(**{key: True}) + + +@pytest.mark.parametrize( + "key", ["stal_enabled", "stal_stats", "stal_min_candidate_guarantee", "stal_zero_positive_rescue"] +) +def test_direct_assigner_rejects_nonboolean_flags(key): + """Direct construction must not accept numeric truth values for boolean experiment switches.""" + with pytest.raises(TypeError, match="must be a boolean"): + TaskAlignedAssigner(**{key: 1}) diff --git a/tests/test_stal_policy_interactions.py b/tests/test_stal_policy_interactions.py new file mode 100644 index 000000000..b62fb0c65 --- /dev/null +++ b/tests/test_stal_policy_interactions.py @@ -0,0 +1,110 @@ +"""Regression coverage for branch budgets and interacting candidate policies.""" + +from types import SimpleNamespace + +import pytest +import torch + +from ultralytics.cfg import get_cfg +from ultralytics.nn.tasks import DetectionModel # noqa: F401 - initialize loss imports +from ultralytics.utils.loss import v8DetectionLoss +from ultralytics.utils.tal import TaskAlignedAssigner + + +@pytest.mark.parametrize("topk,topk2", [(1, None), (7, 1), (10, None)]) +@pytest.mark.parametrize("budget", [10, 15]) +def test_loss_branch_budget(topk, topk2, budget): + """Neutral budgets follow the branch and overrides remain one-to-many only.""" + model = torch.nn.Linear(1, 1) + model.args = get_cfg( + overrides={"stal_candidate_mode": "adaptive", "stal_small_topk": budget, "stal_small_topk_min": budget} + ) + model.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=1, reg_max=16)] + assigner = v8DetectionLoss(model, tal_topk=topk, tal_topk2=topk2).assigner + assert assigner.stal_small_topk == (budget if topk == 10 else topk) + anchors = torch.stack((torch.arange(1.0, 21.0), torch.full((20,), 5.0)), dim=1) + boxes = torch.tensor([[[0.0, 0.0, 22.0, 12.0]]]) + result = assigner( + torch.linspace(0.1, 1.0, 20).view(1, 20, 1), + boxes.expand(1, 20, 4).clone(), + anchors, + torch.zeros(1, 1, 1), + boxes, + torch.ones(1, 1, 1, dtype=torch.bool), + image_size=(800, 800), + epoch=10, + ) + assert result[3].sum() == (budget if topk == 10 else 1) + + +@pytest.mark.parametrize("budget", [3, 10, 15]) +@pytest.mark.parametrize("cap", [1, 2]) +def test_adaptive_topk_preserves_extra_cap(budget, cap): + """Combining strategies must satisfy both upper bounds.""" + anchors = torch.stack((torch.arange(1.0, 15.0), torch.full((14,), 5.0)), dim=1) + boxes = torch.tensor([[[0.0, 0.0, 10.0, 10.0]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + assigner = TaskAlignedAssigner( + topk=10, + num_classes=1, + stal_candidate_mode="adaptive", + stal_relaxation=10.0, + stal_warmup_epochs=0, + stal_area_threshold=0.5, + stal_max_extra_candidates=cap, + stal_small_topk=budget, + stal_small_topk_min=budget, + stride=[1, 1, 1], + ) + assigner.bs = assigner.n_max_boxes = 1 + selected = assigner.get_pos_mask( + torch.linspace(0.1, 0.9, 14).view(1, 14, 1), + boxes.expand(1, 14, 4).clone(), + torch.zeros(1, 1, 1), + boxes, + anchors, + valid, + image_size=(100, 100), + )[0].bool() + base = assigner.select_candidates_in_gts(anchors, boxes, valid, image_size=(100, 100), relaxation_override=0.0) + assert (selected & ~base).sum() == cap + if budget != 10: + assert selected.sum() <= budget + else: + assert (selected & base).sum() == base.sum() + + +@pytest.mark.parametrize("epoch", [0, 1, 5, 10]) +@pytest.mark.parametrize("scale_mode", ["constant", "sqrt_area"]) +def test_crowding_respects_warmup_and_never_enlarges(epoch, scale_mode): + """Crowding caps the area-scaled expansion with the same epoch warmup.""" + x = torch.linspace(-5.0, 15.0, 401) + 0.013 + anchors = torch.stack((x, torch.full_like(x, 5.0)), dim=1) + boxes = torch.tensor([[[0.0, 0.0, 10.0, 10.0], [0.0, 0.0, 10.0, 10.0]]]) + valid = torch.ones(1, 2, 1, dtype=torch.bool) + common = { + "stal_candidate_mode": "adaptive", + "stal_relaxation": 8.0, + "stal_warmup_epochs": 10, + "stal_area_threshold": 0.5, + "stal_relaxation_scale_mode": scale_mode, + } + plain = TaskAlignedAssigner(**common) + crowd = TaskAlignedAssigner(**common, stal_crowding_mode="candidate_overlap", stal_crowded_relaxation=4.0) + baseline = plain.select_candidates_in_gts(anchors, boxes, valid, image_size=(100, 100), epoch=epoch) + actual = crowd.select_candidates_in_gts(anchors, boxes, valid, image_size=(100, 100), epoch=epoch) + assert not (actual & ~baseline).any() + if scale_mode == "constant": + expected = (x > -2.0 * epoch / 10 + 1e-9) & (x < 10.0 + 2.0 * epoch / 10 - 1e-9) + assert torch.equal(actual, expected.view(1, 1, -1).expand_as(actual)) + + +def test_one_to_one_final_filter_keeps_zero_metric_candidate(): + """An invalid zero-score anchor cannot displace a nominated zero-score anchor.""" + assigner = TaskAlignedAssigner(topk=1) + nominations = torch.tensor([[[0.0, 1.0, 1.0]]]) + _, foreground, selected = assigner.select_highest_overlaps( + nominations, torch.ones_like(nominations), 1, torch.zeros_like(nominations) + ) + assert foreground.sum() == 1 + assert not selected[..., 0].any() diff --git a/tests/test_stal_scale_evaluation.py b/tests/test_stal_scale_evaluation.py new file mode 100644 index 000000000..e262fef8d --- /dev/null +++ b/tests/test_stal_scale_evaluation.py @@ -0,0 +1,183 @@ +import json + +import pytest +from PIL import Image + +from scripts.stal.evaluate_scale_ap import build_coco_ground_truth, evaluate_scale_ap, resolve_split_images + + +def test_build_coco_ground_truth_preserves_original_pixel_area(tmp_path): + """COCO size bins must use original-image pixel areas, not resized training coordinates.""" + image_dir = tmp_path / "images" / "val" + label_dir = tmp_path / "labels" / "val" + image_dir.mkdir(parents=True) + label_dir.mkdir(parents=True) + image_path = image_dir / "000001.jpg" + Image.new("RGB", (200, 100)).save(image_path) + (label_dir / "000001.txt").write_text("0 0.5 0.5 0.1 0.2\n", encoding="utf-8") + + coco = build_coco_ground_truth([image_path], {0: "target"}) + + assert coco["images"] == [{"id": 1, "file_name": "000001.jpg", "width": 200, "height": 100}] + assert coco["annotations"][0]["bbox"] == pytest.approx([90.0, 40.0, 20.0, 20.0]) + assert coco["annotations"][0]["area"] == pytest.approx(400.0) + assert coco["annotations"][0]["category_id"] == 1 + + +def test_resolve_split_images_filters_and_sorts_supported_images(tmp_path): + """Directory splits should ignore labels and return a deterministic image order.""" + Image.new("RGB", (8, 8)).save(tmp_path / "b.jpg") + Image.new("RGB", (8, 8)).save(tmp_path / "a.png") + (tmp_path / "notes.txt").write_text("ignore", encoding="utf-8") + + resolved = resolve_split_images(str(tmp_path)) + + assert [path.name for path in resolved] == ["a.png", "b.jpg"] + + +def test_evaluate_scale_ap_reports_perfect_small_object_prediction(tmp_path): + """The external evaluator should expose COCO small AP from an exact synthetic prediction.""" + image_dir = tmp_path / "images" / "val" + label_dir = tmp_path / "labels" / "val" + image_dir.mkdir(parents=True) + label_dir.mkdir(parents=True) + image_path = image_dir / "000001.jpg" + Image.new("RGB", (200, 100)).save(image_path) + (label_dir / "000001.txt").write_text("0 0.5 0.5 0.1 0.2\n", encoding="utf-8") + annotations_path = tmp_path / "annotations.json" + predictions_path = tmp_path / "predictions.json" + annotations_path.write_text(json.dumps(build_coco_ground_truth([image_path], {0: "target"})), encoding="utf-8") + predictions_path.write_text( + json.dumps([{"image_id": 1, "category_id": 1, "bbox": [90.0, 40.0, 20.0, 20.0], "score": 0.99}]), + encoding="utf-8", + ) + + metrics = evaluate_scale_ap(annotations_path, predictions_path) + + assert metrics["AP"] == pytest.approx(1.0) + assert metrics["AP50"] == pytest.approx(1.0) + assert metrics["AP75"] == pytest.approx(1.0) + assert metrics["APs"] == pytest.approx(1.0) + assert metrics["AP50s"] == pytest.approx(1.0) + assert metrics["AR500"] == pytest.approx(1.0) + assert metrics["ARs500"] == pytest.approx(1.0) + + +@pytest.mark.parametrize("side,expected", [(31, "APs"), (32, "APm"), (96, "APl")]) +def test_visdrone_string_ids_and_exact_area_boundaries(tmp_path, side, expected): + """String IDs must work and boundary GTs must belong to exactly one size bin.""" + image_id = "0000001_02999_d_0000005" + bbox = [10, 10, side, side] + gt = { + "images": [{"id": image_id, "file_name": image_id + ".jpg", "width": 200, "height": 200}], + "annotations": [ + {"id": 1, "image_id": image_id, "category_id": 1, "bbox": bbox, "area": side * side, "iscrowd": 0} + ], + "categories": [{"id": 1, "name": "target"}], + } + annotation_path = tmp_path / "gt.json" + prediction_path = tmp_path / "pred.json" + annotation_path.write_text(json.dumps(gt), encoding="utf-8") + prediction_path.write_text( + json.dumps([{"image_id": image_id, "category_id": 1, "bbox": bbox, "score": 0.99}]), encoding="utf-8" + ) + metrics = evaluate_scale_ap(annotation_path, prediction_path) + assert metrics[expected] == pytest.approx(1.0) + for key in {"APs", "APm", "APl"} - {expected}: + assert metrics[key] == -1.0 + assert json.loads(annotation_path.read_text())["images"][0]["id"] == image_id + + +def evaluate_fixture(tmp_path, predictions, annotations=None): + """Evaluate explicit geometry so expected AP does not depend on the implementation under test.""" + gt = { + "images": [{"id": 1, "width": 1000, "height": 1000}], + "categories": [{"id": 1, "name": "target"}], + "annotations": annotations + if annotations is not None + else [{"id": 1, "image_id": 1, "category_id": 1, "bbox": [0, 0, 20, 20], "area": 400, "iscrowd": 0}], + } + annotation_path, prediction_path = tmp_path / "gt.json", tmp_path / "pred.json" + annotation_path.write_text(json.dumps(gt), encoding="utf-8") + prediction_path.write_text(json.dumps(predictions), encoding="utf-8") + return evaluate_scale_ap(annotation_path, prediction_path) + + +def test_aps_averages_iou_thresholds_instead_of_ap50(tmp_path): + """IoU=0.6 passes exactly .50/.55/.60, yielding APs=.3 while AP50s=1.""" + metrics = evaluate_fixture(tmp_path, [{"image_id": 1, "category_id": 1, "bbox": [5, 0, 20, 20], "score": 0.9}]) + assert metrics["APs"] == pytest.approx(0.3) + assert metrics["AP50s"] == pytest.approx(1) + assert metrics["AP75"] == pytest.approx(0) + assert metrics["ARs500"] == pytest.approx(0.3) + + +def test_max_dets_500_keeps_rank_500_and_drops_rank_501(tmp_path): + """A sole TP at rank 500 survives; at rank 501 it cannot contribute.""" + false_positive = {"image_id": 1, "category_id": 1, "bbox": [100, 100, 20, 20], "score": 0.9} + true_positive = {"image_id": 1, "category_id": 1, "bbox": [0, 0, 20, 20], "score": 0.1} + metrics = evaluate_fixture(tmp_path, [false_positive] * 499 + [true_positive]) + assert metrics["APs"] == pytest.approx(1 / 500) + assert metrics["ARs500"] == pytest.approx(1) + metrics = evaluate_fixture(tmp_path, [false_positive] * 500 + [true_positive]) + assert metrics["APs"] == 0 + assert metrics["ARs500"] == 0 + + +def test_no_detections_is_zero_but_absent_gt_bin_is_unavailable(tmp_path): + """A failed detector is not the same as an unavailable metric.""" + metrics = evaluate_fixture(tmp_path, []) + assert metrics["APs"] == metrics["ARs500"] == 0 + assert metrics["APm"] == metrics["APl"] == -1 + assert all(value == -1 for value in evaluate_fixture(tmp_path, [], annotations=[]).values()) + + +@pytest.mark.parametrize( + "update", + [{"category_id": 2}, {"bbox": [0, 0, -1, 20]}, {"score": float("nan")}, {"score": -0.1}, {"score": 1.1}], +) +def test_invalid_predictions_cannot_silently_improve_metrics(tmp_path, update): + """Unknown classes and invalid numbers must not be silently ignored by the external evaluator.""" + prediction = {"image_id": 1, "category_id": 1, "bbox": [0, 0, 20, 20], "score": 0.9} + with pytest.raises(ValueError): + evaluate_fixture(tmp_path, [{**prediction, **update}]) + + +@pytest.mark.parametrize("label", ["0.5 .5 .5 .1 .1", "0 .5 .5 nan .1", "0 .5 .5 -.1 .1"]) +def test_invalid_yolo_labels_are_rejected(tmp_path, label): + """Malformed GT must not change class identity or area silently.""" + images, labels = tmp_path / "images", tmp_path / "labels" + images.mkdir() + labels.mkdir() + image = images / "1.jpg" + Image.new("RGB", (100, 100)).save(image) + (labels / "1.txt").write_text(label, encoding="utf-8") + with pytest.raises(ValueError): + build_coco_ground_truth([image], {0: "target"}) + + +def test_missing_labels_are_not_silently_counted_as_background(tmp_path): + """Incomplete validation mirrors must fail rather than inflate recall by deleting GT.""" + images = tmp_path / "images" + images.mkdir() + image = images / "1.jpg" + Image.new("RGB", (100, 100)).save(image) + with pytest.raises(FileNotFoundError, match="Missing evaluation label"): + build_coco_ground_truth([image], {0: "target"}) + with pytest.raises(ValueError, match="no images"): + build_coco_ground_truth([], {0: "target"}) + + +@pytest.mark.parametrize("update", [{"area": 10000}, {"category_id": 2}, {"bbox": [0, 0, -20, -20]}]) +def test_external_ground_truth_cannot_silently_change_metric_population(tmp_path, update): + """Direct JSON evaluation must enforce bbox area and class validity just like the YOLO converter.""" + annotation = {"id": 1, "image_id": 1, "category_id": 1, "bbox": [0, 0, 20, 20], "area": 400, "iscrowd": 0} + with pytest.raises(ValueError): + evaluate_fixture(tmp_path, [], annotations=[{**annotation, **update}]) + + +def test_duplicate_annotation_ids_are_rejected_before_coco_indexing(tmp_path): + """COCO's ID index must not overwrite distinct GT records sharing an annotation ID.""" + annotation = {"id": 1, "image_id": 1, "category_id": 1, "bbox": [0, 0, 20, 20], "area": 400, "iscrowd": 0} + with pytest.raises(ValueError, match="Duplicate annotation"): + evaluate_fixture(tmp_path, [], annotations=[annotation, {**annotation, "bbox": [50, 50, 20, 20]}]) diff --git a/tests/test_stal_simd.py b/tests/test_stal_simd.py new file mode 100644 index 000000000..97073569c --- /dev/null +++ b/tests/test_stal_simd.py @@ -0,0 +1,64 @@ +"""Verify the opt-in SimD-inspired small-target TAL matching adaptation.""" + +import math + +import pytest +import torch + +from ultralytics.cfg import check_cfg, get_cfg +from ultralytics.utils.tal import TaskAlignedAssigner + + +def test_simd_matches_official_visdrone_formula_constants(): + """Match the official implementation's x=6.13, y=4.59 location-and-shape formula.""" + gt = torch.tensor([[0.0, 0.0, 4.0, 2.0]]) + pred = torch.tensor([[1.0, 0.0, 7.0, 4.0]]) + gt_xywh = torch.tensor([2.0, 1.0, 4.0, 2.0]) + pred_xywh = torch.tensor([4.0, 2.0, 6.0, 4.0]) + scale = (gt_xywh[2:] + pred_xywh[2:]) / torch.tensor([6.13, 4.59]) + location = (((gt_xywh[:2] - pred_xywh[:2]) / scale).square().sum()).sqrt() + shape = (((gt_xywh[2:] - pred_xywh[2:]) / scale).square().sum()).sqrt() + expected = math.exp(-float(location + shape)) + assert TaskAlignedAssigner.simd_similarity(gt, pred).item() == pytest.approx(expected) + + +def test_simd_changes_only_small_target_matching_and_keeps_ciou_targets(): + """Use SimD for small-target ranking while decoupling final soft targets to CIoU.""" + assigner = TaskAlignedAssigner( + num_classes=1, + topk=2, + stal_candidate_mode="pure", + stal_area_threshold=0.01, + stal_simd_weight=1.0, + stal_nwd_target_mode="ciou", + ) + scores = torch.tensor([[[0.5], [0.5]]]) + predicted = torch.tensor([[[0.0, 0.0, 4.0, 4.0], [1.0, 0.0, 5.0, 4.0]]]) + anchors = torch.tensor([[2.0, 2.0], [3.0, 2.0]]) + labels = torch.zeros(1, 1, 1) + boxes = torch.tensor([[[0.0, 0.0, 4.0, 4.0]]]) + valid = torch.ones(1, 1, 1, dtype=torch.bool) + assigner.bs = 1 + assigner.n_max_boxes = 1 + mask, _, matching, _, targets = assigner.get_pos_mask( + scores, predicted, labels, boxes, anchors, valid, image_size=(100, 100) + ) + assert mask.any() + assert matching[0, 0, 1] != pytest.approx(targets[0, 0, 1]) + assert targets[0, 0, 0] == pytest.approx(1.0) + + +@pytest.mark.parametrize("value", [-0.1, 1.1, float("nan")]) +def test_simd_rejects_invalid_weights(value): + """Reject invalid weights at both public configuration boundaries.""" + with pytest.raises(ValueError): + TaskAlignedAssigner(stal_simd_weight=value) + with pytest.raises(ValueError): + check_cfg({"stal_simd_weight": value}) + + +def test_simd_default_and_type_contract(): + """Default is disabled and programmatic strings cannot bypass type checking.""" + assert get_cfg().stal_simd_weight == 0.0 + with pytest.raises(TypeError): + check_cfg({"stal_simd_weight": "1.0"}) diff --git a/tests/test_stal_telemetry_contract.py b/tests/test_stal_telemetry_contract.py new file mode 100644 index 000000000..26ebfd14b --- /dev/null +++ b/tests/test_stal_telemetry_contract.py @@ -0,0 +1,112 @@ +"""Keep relative STAL and absolute small-target telemetry populations distinct.""" + +from types import SimpleNamespace + +import pytest +import torch + +from ultralytics.cfg import get_cfg +from ultralytics.engine.trainer import BaseTrainer +from ultralytics.nn.tasks import DetectionModel # noqa: F401 +from ultralytics.utils.loss import E2ELoss, v8DetectionLoss +from ultralytics.utils.tal import TaskAlignedAssigner + + +@pytest.mark.parametrize("side", [24, 32, 40, 80]) +def test_score_groups_match_their_gt_area_definitions(side): + model = torch.nn.Linear(1, 1) + model.args = get_cfg(overrides={"stal_stats": True}) + model.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=1, reg_max=16)] + loss = v8DetectionLoss(model) + predictions = { + "feats": [torch.zeros(1, 1, s, s) for s in (100, 50, 25)], + "boxes": torch.zeros(1, 64, 13125), + "scores": torch.zeros(1, 1, 13125), + } + batch = { + "batch_idx": torch.zeros(1), + "cls": torch.zeros(1, 1), + "bboxes": torch.tensor([[0.5, 0.5, side / 800, side / 800]]), + "epoch": 10, + } + loss.get_assigned_targets_and_loss(predictions, batch) + score = loss.pop_target_score_stats() + assignment = loss.pop_assignment_stats() + assert score.shape == (9,) # STAL, all, COCO-style small; three values each + assert score[3] > 0 + assert score[0] == assignment[1] + assert score[6] == assignment[4] + assert bool(score[0]) == (side < 80) + assert bool(score[6]) == (side < 32) + assert not loss.pop_target_score_stats().any() + + +@pytest.mark.parametrize("constructor", [TaskAlignedAssigner, lambda **kw: get_cfg(overrides=kw)]) +def test_guarantee_cannot_override_a_configured_capacity_floor(constructor): + with pytest.raises(ValueError, match="guarantee.*floor"): + constructor(stal_candidate_mode="adaptive", stal_min_candidate_guarantee=True, stal_candidate_iou_floor=0.75) + + +def test_epoch_attempt_discards_all_abandoned_statistics(): + model = torch.nn.Linear(1, 1) + model.args = get_cfg(overrides={"stal_stats": True}) + model.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=1, reg_max=16)] + criterion = v8DetectionLoss(model) + model.criterion = criterion + counters = [ + criterion._assignment_stats, + criterion._rescue_stats, + criterion._assignment_stage_stats, + criterion._target_score_stats, + ] + for counter in counters: + counter.fill_(7) + BaseTrainer._reset_assignment_metric_state(SimpleNamespace(model=model)) + assert all(not counter.any() for counter in counters) + criterion._assignment_stats[0] = 3 + assert criterion.pop_assignment_stats()[0] == 3 + + +def test_initial_epoch_reset_accepts_lazy_criterion(): + BaseTrainer._reset_assignment_metric_state(SimpleNamespace(model=torch.nn.Linear(1, 1))) + + +@pytest.mark.parametrize( + "name,buffer", + [ + ("pop_assignment_stats", "_assignment_stats"), + ("pop_rescue_stats", "_rescue_stats"), + ("pop_assignment_stage_stats", "_assignment_stage_stats"), + ("pop_target_score_stats", "_target_score_stats"), + ], +) +def test_e2e_reports_one_to_many_without_double_counting_and_drains_both(name, buffer): + model = torch.nn.Linear(1, 1) + model.args = get_cfg(overrides={"stal_stats": True}) + model.model = [SimpleNamespace(stride=torch.tensor([8.0, 16.0, 32.0]), nc=1, reg_max=16)] + loss = E2ELoss(model) + getattr(loss.one2many, buffer).fill_(3) + getattr(loss.one2one, buffer).fill_(7) + assert (getattr(loss, name)() == 3).all() + assert not getattr(loss.one2many, buffer).any() + assert not getattr(loss.one2one, buffer).any() + + +def test_new_stats_cannot_append_to_historical_csv_schema(tmp_path): + path = tmp_path / "results.csv" + previous = "epoch,time,assign/small_nonzero_score_pos\n1,1,3\n" + path.write_text(previous, encoding="utf-8") + trainer = SimpleNamespace(csv=path, train_time_start=0, epoch=1) + metrics = {"assign/stal_nonzero_score_pos": 3, "assign/small_nonzero_score_pos": 1} + with pytest.raises(ValueError, match="schema changed"): + BaseTrainer.save_metrics(trainer, metrics) + assert path.read_text() == previous + + +def test_new_stats_append_consistently_in_a_fresh_run(tmp_path): + trainer = SimpleNamespace(csv=tmp_path / "results.csv", train_time_start=0, epoch=0) + metrics = {"assign/stal_nonzero_score_pos": 3, "assign/small_nonzero_score_pos": 1} + BaseTrainer.save_metrics(trainer, metrics) + trainer.epoch = 1 + BaseTrainer.save_metrics(trainer, metrics) + assert len(trainer.csv.read_text().splitlines()) == 3 diff --git a/tests/test_stal_visdrone_export.py b/tests/test_stal_visdrone_export.py new file mode 100644 index 000000000..fe318e799 --- /dev/null +++ b/tests/test_stal_visdrone_export.py @@ -0,0 +1,114 @@ +import json + +import pytest +from PIL import Image + +from scripts.stal.export_visdrone_results import export_visdrone_results + + +def test_export_visdrone_results_writes_official_rows_and_empty_images(tmp_path): + """Exporter should preserve all images and emit the official eight-field DET row layout.""" + images = tmp_path / "images" + images.mkdir() + Image.new("RGB", (20, 10)).save(images / "000001.jpg") + Image.new("RGB", (20, 10)).save(images / "000002.jpg") + predictions = tmp_path / "predictions.json" + predictions.write_text( + json.dumps( + [ + { + "image_id": 1, + "file_name": "000001.jpg", + "category_id": 4, + "bbox": [1.0, 2.0, 3.0, 4.0], + "score": 0.875, + } + ] + ), + encoding="utf-8", + ) + + summary = export_visdrone_results(predictions, images, tmp_path / "results") + + assert summary == {"images": 2, "detections": 1, "empty_images": 1} + assert (tmp_path / "results" / "000001.txt").read_text(encoding="utf-8") == ( + "1.000,2.000,3.000,4.000,0.87500000,4,-1,-1\n" + ) + assert (tmp_path / "results" / "000002.txt").read_text(encoding="utf-8") == "" + + +@pytest.mark.parametrize("category", [0, 11, 1.5, float("nan")]) +def test_export_visdrone_results_rejects_non_evaluated_categories(tmp_path, category): + """Ignored-region and others categories must not enter official detector submissions.""" + images = tmp_path / "images" + images.mkdir() + Image.new("RGB", (20, 10)).save(images / "000001.jpg") + predictions = tmp_path / "predictions.json" + predictions.write_text( + json.dumps([{"image_id": 1, "category_id": category, "bbox": [1, 2, 3, 4], "score": 0.5}]), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="category_id"): + export_visdrone_results(predictions, images, tmp_path / "results") + + +@pytest.mark.parametrize("update", [{"bbox": [0, 0, 0, 2]}, {"score": float("inf")}, {"score": -0.1}, {"score": 1.1}]) +def test_invalid_export_does_not_write_result_files(tmp_path, update): + """Reject malformed values before creating a partial official submission.""" + images = tmp_path / "images" + images.mkdir() + Image.new("RGB", (20, 20)).save(images / "1.jpg") + predictions = tmp_path / "pred.json" + prediction = {"image_id": 1, "category_id": 1, "bbox": [0, 0, 2, 2], "score": 0.9} + predictions.write_text(json.dumps([{**prediction, **update}]), encoding="utf-8") + with pytest.raises(ValueError): + export_visdrone_results(predictions, images, tmp_path / "results") + assert not (tmp_path / "results").exists() + + +@pytest.mark.parametrize("file_name", ["000002.jpg", "missing.jpg"]) +def test_export_rejects_conflicting_image_identifiers(tmp_path, file_name): + """A prediction must not move to another image or hide a bad filename behind its numeric ID.""" + images = tmp_path / "images" + images.mkdir() + for name in ("000001.jpg", "000002.jpg"): + Image.new("RGB", (20, 20)).save(images / name) + predictions = tmp_path / "pred.json" + predictions.write_text( + json.dumps([{"image_id": 1, "file_name": file_name, "category_id": 1, "bbox": [0, 0, 2, 2], "score": 0.9}]), + encoding="utf-8", + ) + with pytest.raises(ValueError): + export_visdrone_results(predictions, images, tmp_path / "results") + assert not (tmp_path / "results").exists() + + +@pytest.mark.parametrize("names", [("1.jpg", "1.png"), ("01.jpg", "001.jpg")]) +def test_export_rejects_colliding_image_names(tmp_path, names): + """Neither output TXT names nor numeric image IDs may alias two images.""" + images = tmp_path / "images" + images.mkdir() + for name in names: + Image.new("RGB", (20, 20)).save(images / name) + predictions = tmp_path / "pred.json" + predictions.write_text("[]", encoding="utf-8") + with pytest.raises(ValueError): + export_visdrone_results(predictions, images, tmp_path / "results") + + +def test_export_rejects_stale_result_files_without_overwriting(tmp_path): + """Reusing a directory from another split must not leave a mixed submission behind.""" + images = tmp_path / "images" + images.mkdir() + Image.new("RGB", (20, 20)).save(images / "1.jpg") + predictions = tmp_path / "pred.json" + predictions.write_text("[]", encoding="utf-8") + output = tmp_path / "results" + output.mkdir() + (output / "old.txt").write_text("old run", encoding="utf-8") + (output / "1.txt").write_text("current run", encoding="utf-8") + with pytest.raises(ValueError, match="Unexpected result"): + export_visdrone_results(predictions, images, output) + assert (output / "1.txt").read_text() == "current run" + assert (output / "old.txt").read_text() == "old run" diff --git a/tests/test_tal_mps_regression.py b/tests/test_tal_mps_regression.py index 39b897cbb..49f90b78a 100644 --- a/tests/test_tal_mps_regression.py +++ b/tests/test_tal_mps_regression.py @@ -18,7 +18,9 @@ def test_task_aligned_assigner_box_metrics_match_expanded_reference(): assigner = TaskAlignedAssigner(num_classes=classes) assigner.bs, assigner.n_max_boxes = batch_size, max_boxes - align_metric, overlaps = assigner.get_box_metrics(pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt) + align_metric, overlaps, target_overlaps = assigner.get_box_metrics( + pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt + ) reference_scores = torch.zeros(batch_size, max_boxes, anchors) reference_overlaps = torch.zeros_like(reference_scores) @@ -31,6 +33,7 @@ def test_task_aligned_assigner_box_metrics_match_expanded_reference(): reference_align_metric = reference_scores.pow(assigner.alpha) * reference_overlaps.pow(assigner.beta) torch.testing.assert_close(overlaps, reference_overlaps) + torch.testing.assert_close(target_overlaps, reference_overlaps) torch.testing.assert_close(align_metric, reference_align_metric) diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index de72fb6fd..9f9db0076 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 @@ -267,6 +268,13 @@ CFG_FLOAT_KEYS = frozenset( { # integer or float arguments, i.e. x=2 and x=2.0 "warmup_epochs", + "stal_relaxation", + "stal_relaxation_min", + "stal_min_size_stride_ratio", + "stal_warmup_epochs", + "stal_crowded_relaxation", + "stal_rescue_floor_decay_epochs", + "stal_nwd_constant", "box", "cls", "dfl", @@ -317,6 +325,15 @@ "iou", "fraction", "multi_scale", + "stal_area_threshold", + "stal_candidate_iou_floor", + "stal_expanded_quality_ratio", + "stal_expanded_score_floor", + "stal_rescue_score_floor", + "stal_nwd_weight", + "stal_nwd_target_weight", + "stal_nwd_zero_score_floor", + "stal_simd_weight", } ) MIXTURE_INT_KEYS = frozenset( @@ -372,6 +389,10 @@ "line_width", "nbs", "save_period", + "stal_small_topk", + "stal_small_topk_min", + "stal_min_base_candidates", + "stal_max_extra_candidates", } ) | MIXTURE_INT_KEYS CFG_INT_MIN = { # minimum valid values for integer arguments used as divisors, sizes or seeds @@ -380,6 +401,9 @@ "mask_ratio": 1, "vid_stride": 1, "seed": 0, + "stal_small_topk": 1, + "stal_small_topk_min": 1, + "stal_min_base_candidates": 0, "moe_prune_calibration_steps": 1, "mot_sparse_train_warmup_steps": 0, "mot_local_attn_window": 0, @@ -474,6 +498,10 @@ "foundation_semantic_distill", "foundation_multitask", "foundation_multitask_enabled", + "stal_enabled", + "stal_min_candidate_guarantee", + "stal_stats", + "stal_zero_positive_rescue", } ) | MIXTURE_BOOL_KEYS MIXTURE_STR_KEYS = frozenset( @@ -499,6 +527,10 @@ { "optimizer", "split", + "stal_candidate_mode", + "stal_crowding_mode", + "stal_relaxation_scale_mode", + "stal_nwd_target_mode", "copy_paste_mode", "auto_augment", "foundation_teacher", @@ -519,6 +551,7 @@ FOUNDATION_RELATION_MODES = frozenset({"sampled", "full"}) FOUNDATION_DTYPES = frozenset({"auto", "fp32", "fp16", "bf16"}) FOUNDATION_TARGET_LEVELS = frozenset({"p3", "p4", "p5"}) +STAL_CANDIDATE_MODES = frozenset({"pure", "fixed", "adaptive"}) # fmt: on LORA_RUNTIME_METADATA_KEYS = frozenset( { @@ -749,6 +782,106 @@ def check_cfg(cfg: dict, hard: bool = True) -> None: else: cfg[k] = scheme + validate_stal_config(cfg) + + +def validate_stal_config(cfg: dict) -> None: + """Validate STAL candidate-policy values and backward-compatible parameter relationships.""" + for key, value in cfg.items(): + if key.startswith("stal_") and key in CFG_FLOAT_KEYS | CFG_FRACTION_KEYS | CFG_INT_KEYS: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"'{key}' must be numeric, not {type(value).__name__}.") + if not math.isfinite(value): + raise ValueError(f"'{key}' must be finite.") + + mode = cfg.get("stal_candidate_mode", DEFAULT_CFG_DICT["stal_candidate_mode"]) + enabled = cfg.get("stal_enabled", DEFAULT_CFG_DICT["stal_enabled"]) + threshold = cfg.get("stal_area_threshold", DEFAULT_CFG_DICT["stal_area_threshold"]) + small_topk = cfg.get("stal_small_topk", DEFAULT_CFG_DICT["stal_small_topk"]) + small_topk_min = cfg.get("stal_small_topk_min", DEFAULT_CFG_DICT["stal_small_topk_min"]) + min_base_candidates = cfg.get("stal_min_base_candidates", DEFAULT_CFG_DICT["stal_min_base_candidates"]) + max_extra_candidates = cfg.get("stal_max_extra_candidates", DEFAULT_CFG_DICT["stal_max_extra_candidates"]) + crowding_mode = cfg.get("stal_crowding_mode", DEFAULT_CFG_DICT["stal_crowding_mode"]) + crowded_relaxation = cfg.get("stal_crowded_relaxation", DEFAULT_CFG_DICT["stal_crowded_relaxation"]) + candidate_floor = cfg.get("stal_candidate_iou_floor", DEFAULT_CFG_DICT["stal_candidate_iou_floor"]) + if ( + cfg.get("stal_min_candidate_guarantee", DEFAULT_CFG_DICT["stal_min_candidate_guarantee"]) + and candidate_floor > 0 + ): + raise ValueError("stal_min_candidate_guarantee cannot be combined with a positive stal_candidate_iou_floor") + expanded_quality_ratio = cfg.get("stal_expanded_quality_ratio", DEFAULT_CFG_DICT["stal_expanded_quality_ratio"]) + expanded_score_floor = cfg.get("stal_expanded_score_floor", DEFAULT_CFG_DICT["stal_expanded_score_floor"]) + relaxation = cfg.get("stal_relaxation", DEFAULT_CFG_DICT["stal_relaxation"]) + relaxation_min = cfg.get("stal_relaxation_min", DEFAULT_CFG_DICT["stal_relaxation_min"]) + min_size_stride_ratio = cfg.get("stal_min_size_stride_ratio", DEFAULT_CFG_DICT["stal_min_size_stride_ratio"]) + relaxation_scale_mode = cfg.get("stal_relaxation_scale_mode", DEFAULT_CFG_DICT["stal_relaxation_scale_mode"]) + warmup = cfg.get("stal_warmup_epochs", DEFAULT_CFG_DICT["stal_warmup_epochs"]) + rescue = cfg.get("stal_zero_positive_rescue", DEFAULT_CFG_DICT["stal_zero_positive_rescue"]) + rescue_floor = cfg.get("stal_rescue_score_floor", DEFAULT_CFG_DICT["stal_rescue_score_floor"]) + rescue_decay = cfg.get("stal_rescue_floor_decay_epochs", DEFAULT_CFG_DICT["stal_rescue_floor_decay_epochs"]) + nwd_weight = cfg.get("stal_nwd_weight", DEFAULT_CFG_DICT["stal_nwd_weight"]) + nwd_constant = cfg.get("stal_nwd_constant", DEFAULT_CFG_DICT["stal_nwd_constant"]) + nwd_target_mode = cfg.get("stal_nwd_target_mode", DEFAULT_CFG_DICT["stal_nwd_target_mode"]) + nwd_target_weight = cfg.get("stal_nwd_target_weight", DEFAULT_CFG_DICT["stal_nwd_target_weight"]) + nwd_zero_score_floor = cfg.get("stal_nwd_zero_score_floor", DEFAULT_CFG_DICT["stal_nwd_zero_score_floor"]) + simd_weight = cfg.get("stal_simd_weight", DEFAULT_CFG_DICT["stal_simd_weight"]) + if mode not in STAL_CANDIDATE_MODES: + raise ValueError(f"'stal_candidate_mode={mode}' is invalid. Use one of {sorted(STAL_CANDIDATE_MODES)}.") + if not 0.0 < threshold <= 1.0: + raise ValueError(f"'stal_area_threshold={threshold}' is invalid. Use a value in (0, 1].") + if small_topk < 1: + raise ValueError("'stal_small_topk' must be at least 1.") + if not 1 <= small_topk_min <= small_topk: + raise ValueError("'stal_small_topk_min' must be in [1, stal_small_topk].") + if min_base_candidates < 0: + raise ValueError("'stal_min_base_candidates' must be non-negative.") + if max_extra_candidates < 0: + raise ValueError("'stal_max_extra_candidates' must be non-negative.") + if crowding_mode not in {"none", "candidate_overlap"}: + raise ValueError("'stal_crowding_mode' must be 'none' or 'candidate_overlap'.") + if not 0.0 <= crowded_relaxation <= relaxation: + raise ValueError("'stal_crowded_relaxation' must be in [0, stal_relaxation].") + if not 0.0 <= candidate_floor <= 1.0: + raise ValueError("'stal_candidate_iou_floor' must be in [0, 1].") + if not 0.0 <= expanded_quality_ratio <= 1.0: + raise ValueError("'stal_expanded_quality_ratio' must be in [0, 1].") + if not 0.0 <= expanded_score_floor <= 1.0: + raise ValueError("'stal_expanded_score_floor' must be in [0, 1].") + if relaxation < 0.0: + raise ValueError(f"'stal_relaxation={relaxation}' must be non-negative.") + if relaxation_scale_mode not in {"constant", "sqrt_area"}: + raise ValueError("'stal_relaxation_scale_mode' must be 'constant' or 'sqrt_area'.") + if not 0.0 <= relaxation_min <= relaxation: + raise ValueError("'stal_relaxation_min' must be in [0, stal_relaxation].") + if min_size_stride_ratio < 0.0: + raise ValueError("'stal_min_size_stride_ratio' must be non-negative.") + if min_size_stride_ratio > 0.0 and relaxation > 0.0: + raise ValueError("'stal_min_size_stride_ratio>0' requires 'stal_relaxation=0' for an isolated strategy.") + if min_size_stride_ratio > 0.0 and relaxation_scale_mode != "constant": + raise ValueError("'stal_min_size_stride_ratio>0' requires 'stal_relaxation_scale_mode=constant'.") + if warmup < 0.0: + raise ValueError(f"'stal_warmup_epochs={warmup}' must be non-negative.") + if not 0.0 <= rescue_floor <= 1.0: + raise ValueError(f"'stal_rescue_score_floor={rescue_floor}' is invalid. Use a value in [0, 1].") + if rescue_decay < 0.0: + raise ValueError(f"'stal_rescue_floor_decay_epochs={rescue_decay}' must be non-negative.") + if not 0.0 <= nwd_weight <= 1.0: + raise ValueError(f"'stal_nwd_weight={nwd_weight}' is invalid. Use a value in [0, 1].") + if nwd_constant <= 0.0: + raise ValueError(f"'stal_nwd_constant={nwd_constant}' must be positive.") + if nwd_target_mode not in {"match", "ciou", "weighted"}: + raise ValueError("'stal_nwd_target_mode' must be 'match', 'ciou', or 'weighted'.") + if not 0.0 <= nwd_target_weight <= nwd_weight: + raise ValueError("'stal_nwd_target_weight' must be in [0, stal_nwd_weight].") + if not 0.0 <= nwd_zero_score_floor <= 1.0: + raise ValueError("'stal_nwd_zero_score_floor' must be in [0, 1].") + if not 0.0 <= simd_weight <= 1.0: + raise ValueError("'stal_simd_weight' must be in [0, 1].") + if enabled and mode == "pure": + raise ValueError("'stal_enabled=True' conflicts with 'stal_candidate_mode=pure'.") + if rescue_floor > 0.0 and not rescue: + raise ValueError("'stal_rescue_score_floor>0' requires 'stal_zero_positive_rescue=True'.") + def _foundation_transformers_available() -> bool: """Return whether the optional Transformers package can be discovered without importing it.""" diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml index ac9aa1572..4d233175d 100644 --- a/ultralytics/cfg/default.yaml +++ b/ultralytics/cfg/default.yaml @@ -104,6 +104,34 @@ weight_decay: 0.0005 # (float) weight decay (L2 regularization) warmup_epochs: 3.0 # (float) warmup epochs (fractions allowed) warmup_momentum: 0.8 # (float) initial momentum during warmup warmup_bias_lr: 0.1 # (float) bias learning rate during warmup +stal_enabled: False # (bool) enable small-target area-aware label assignment for detection training +stal_stats: False # (bool) record per-epoch positive assignments for STAL-target and all-target groups +stal_candidate_mode: fixed # (str) candidate policy: pure TAL, existing fixed-stride, or adaptive STAL +stal_candidate_iou_floor: 0.0 # (float) minimum geometric IoU capacity of expanded small-GT candidates; 0 disables +stal_expanded_quality_ratio: 0.0 # (float) minimum expanded-candidate alignment relative to best base candidate; 0 disables +stal_expanded_score_floor: 1.0 # (float) minimum target-score multiplier for adaptive-only candidates; 1 disables +stal_max_extra_candidates: 0 # (int) max adaptive-only additions per small GT after preserving base TAL; 0 disables +stal_area_threshold: 0.01 # (float) target-to-image area ratio below which STAL is applied +stal_small_topk: 10 # (int) adaptive-mode top-k for small targets; 10 preserves standard TAL +stal_small_topk_min: 10 # (int) minimum top-k for the tiniest targets; below stal_small_topk enables area adaptation +stal_min_base_candidates: 0 # (int) expand only GTs with fewer base candidates; 0 expands every small GT +stal_min_candidate_guarantee: False # (bool) add nearest point only when an adaptive small GT has no candidate +stal_crowding_mode: none # (str) crowded-small-target policy: none or candidate_overlap +stal_crowded_relaxation: 0.0 # (float) total width/height increase for crowded small GTs +stal_relaxation: 8.0 # (float) maximum total width/height added to the candidate region in pixels +stal_relaxation_scale_mode: constant # (str) constant or sqrt_area interpolation by augmented target area +stal_relaxation_min: 0.0 # (float) expansion near the small-area threshold in sqrt_area mode +stal_min_size_stride_ratio: 0.0 # (float) minimum candidate width/height in multiples of minimum stride; 0 disables +stal_warmup_epochs: 10.0 # (float) epochs used to linearly ramp STAL relaxation +stal_zero_positive_rescue: False # (bool) give an uncovered small GT its best positive-quality legal candidate +stal_rescue_score_floor: 0.0 # (float) bootstrap supervision floor for rescued zero-quality candidates; 0 disables +stal_rescue_floor_decay_epochs: 0.0 # (float) epochs used to linearly decay the bootstrap floor to zero; 0 keeps it +stal_nwd_weight: 0.0 # (float) CIoU-to-NWD blend weight for small-target assignment quality; 0 preserves existing TAL +stal_nwd_constant: 12.8 # (float) positive normalization constant in exp(-sqrt(W2)/C) NWD similarity +stal_nwd_target_mode: match # (str) target-score quality: match uses NWD blend, ciou decouples supervision strength +stal_nwd_target_weight: 0.0 # (float) NWD share for target scores when target mode is weighted +stal_nwd_zero_score_floor: 0.0 # (float) target-score floor only for selected small anchors with zero CIoU score +stal_simd_weight: 0.0 # (float) CIoU-to-SimD blend for small-target matching; 0 preserves existing behavior distill_model: # (str, optional) path to teacher model for knowledge distillation dis: 6.0 # (float) distillation loss weight diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py index bc41feb69..9a5a84511 100644 --- a/ultralytics/engine/trainer.py +++ b/ultralytics/engine/trainer.py @@ -716,6 +716,7 @@ def _do_train(self): self._model_train() self._reset_foundation_metric_state() + self._reset_assignment_metric_state() if not getattr(self.train_loader, "set_epoch", lambda _: False)(epoch) and RANK != -1: self.train_loader.sampler.set_epoch(epoch) pbar = enumerate(self.train_loader) @@ -757,6 +758,12 @@ def _do_train(self): with sync_context: with autocast(self.amp): batch = self.preprocess_batch(batch) + if ( + self.args.stal_enabled + or self.args.stal_stats + or self.args.stal_candidate_mode == "adaptive" + ): + batch["epoch"] = epoch if self.args.compile: # Decouple inference and loss calculations for improved compile performance preds = self.model(batch["img"]) @@ -844,6 +851,83 @@ def _do_train(self): if hasattr(unwrap_model(self.model).criterion, "update"): unwrap_model(self.model).criterion.update() + assignment_metrics = {} + pop_assignment_stats = getattr(unwrap_model(self.model).criterion, "pop_assignment_stats", None) + if callable(pop_assignment_stats) and self.args.stal_stats: + assignment_counts = pop_assignment_stats() + if RANK != -1: + dist.all_reduce(assignment_counts, op=dist.ReduceOp.SUM) + groups = ("stal", "small", "medium", "large", "all") + assignment_metrics = {} + for index, group in enumerate(groups): + gt, pos, zero = (int(x.item()) for x in assignment_counts[index * 3 : index * 3 + 3]) + suffix = f"{group}_gt" + assignment_metrics[f"assign/{suffix}"] = gt + pos_suffix = "stal_pos" if group == "stal" else "all_pos" if group == "all" else f"{group}_pos" + assignment_metrics[f"assign/{pos_suffix}"] = pos + mean_suffix = ( + "pos_per_stal_gt" + if group == "stal" + else "pos_per_gt" + if group == "all" + else f"pos_per_{group}_gt" + ) + assignment_metrics[f"assign/{mean_suffix}"] = pos / max(gt, 1) + assignment_metrics[f"assign/{group}_zero_pos_gt"] = zero + assignment_metrics[f"assign/{group}_zero_pos_ratio"] = zero / max(gt, 1) + + pop_rescue_stats = getattr(unwrap_model(self.model).criterion, "pop_rescue_stats", None) + if callable(pop_rescue_stats): + rescue_counts = pop_rescue_stats() + if RANK != -1: + dist.all_reduce(rescue_counts, op=dist.ReduceOp.SUM) + rescue_names = ( + "attempted", + "has_legal_candidate", + "has_free_candidate", + "has_positive_quality_candidate", + "proposed", + "succeeded", + "lost_to_conflict", + "bootstrap_proposed", + "bootstrap_succeeded", + ) + assignment_metrics.update( + {f"assign/rescue_{name}": int(value.item()) for name, value in zip(rescue_names, rescue_counts)} + ) + pop_stage_stats = getattr(unwrap_model(self.model).criterion, "pop_assignment_stage_stats", None) + if callable(pop_stage_stats): + stage_counts = pop_stage_stats() + if RANK != -1: + dist.all_reduce(stage_counts, op=dist.ReduceOp.SUM) + stage_names = ( + "stal_stage_gt", + "stal_no_legal_candidate", + "stal_legal_zero_alignment", + "stal_sub_eps_alignment", + "stal_above_eps_alignment", + "stal_topk_missed_nonzero", + "stal_preconflict_positive", + "stal_conflict_lost", + "stal_postconflict_zero", + ) + assignment_metrics.update( + {f"assign/{name}": int(value.item()) for name, value in zip(stage_names, stage_counts)} + ) + pop_target_score_stats = getattr(unwrap_model(self.model).criterion, "pop_target_score_stats", None) + if callable(pop_target_score_stats): + score_stats = pop_target_score_stats() + if RANK != -1: + dist.all_reduce(score_stats, op=dist.ReduceOp.SUM) + for index, group in enumerate(("stal", "all", "small")): + foreground, nonzero, score_sum = ( + float(x.item()) for x in score_stats[index * 3 : index * 3 + 3] + ) + assignment_metrics[f"assign/{group}_nonzero_score_pos"] = nonzero + assignment_metrics[f"assign/{group}_nonzero_score_ratio"] = nonzero / max(foreground, 1.0) + assignment_metrics[f"assign/{group}_target_score_sum"] = score_sum + assignment_metrics[f"assign/{group}_target_score_per_pos"] = score_sum / max(foreground, 1.0) + self.lr = {f"lr/pg{ir}": x["lr"] for ir, x in enumerate(self.optimizer.param_groups)} # for loggers self.run_callbacks("on_train_epoch_end") @@ -877,7 +961,9 @@ def _do_train(self): foundation_metrics = self._mean_foundation_metrics(prefix="train/") if foundation_metrics: self.metrics = {**(self.metrics or {}), **foundation_metrics} - self.save_metrics(metrics={**self.label_loss_items(self.tloss), **self.metrics, **self.lr}) + self.save_metrics( + metrics={**self.label_loss_items(self.tloss), **self.metrics, **assignment_metrics, **self.lr} + ) self.stop |= self.stopper(epoch + 1, self.fitness) or final_epoch if self.args.time: self.stop |= (time.time() - self.train_time_start) > (self.args.time * 3600) @@ -1311,6 +1397,19 @@ def _reset_foundation_metric_state(self) -> None: self.foundation_metric_steps = 0 self.foundation_metric_latest = {} + def _reset_assignment_metric_state(self) -> None: + """Discard counters from an abandoned epoch attempt before replaying batches.""" + criterion = getattr(unwrap_model(self.model), "criterion", None) + for name in ( + "pop_assignment_stats", + "pop_rescue_stats", + "pop_assignment_stage_stats", + "pop_target_score_stats", + ): + reset = getattr(criterion, name, None) + if callable(reset): + reset() + def _collect_foundation_metrics(self) -> None: """Collect the latest Foundation wrapper metrics after a train forward pass.""" model = unwrap_model(self.model) @@ -1366,6 +1465,11 @@ def save_metrics(self, metrics): n = len(metrics) + 2 # number of cols t = time.time() - self.train_time_start self.csv.parent.mkdir(parents=True, exist_ok=True) # ensure parent directory exists + if "assign/stal_nonzero_score_pos" in metrics and self.csv.exists(): + with self.csv.open(encoding="utf-8") as existing: + header = existing.readline().strip().split(",") + if header != ["epoch", "time", *keys]: + raise ValueError("STAL metric schema changed; use a new run directory or the original frozen source") s = "" if self.csv.exists() else ("%s," * n % ("epoch", "time", *keys)).rstrip(",") + "\n" with open(self.csv, "a", encoding="utf-8") as f: f.write(s + ("%.6g," * n % (self.epoch + 1, t, *vals)).rstrip(",") + "\n") diff --git a/ultralytics/utils/loss.py b/ultralytics/utils/loss.py index da08e7be0..bb13d916f 100644 --- a/ultralytics/utils/loss.py +++ b/ultralytics/utils/loss.py @@ -360,6 +360,11 @@ def __init__( self.no = m.nc + m.reg_max * 4 self.reg_max = m.reg_max self.device = device + self.stal_stats = getattr(h, "stal_stats", False) + self._assignment_stats = torch.zeros(15, dtype=torch.long, device=device) + self._rescue_stats = torch.zeros(9, dtype=torch.long, device=device) + self._assignment_stage_stats = torch.zeros(9, dtype=torch.long, device=device) + self._target_score_stats = torch.zeros(9, dtype=torch.float64, device=device) self.use_dfl = m.reg_max > 1 @@ -368,6 +373,12 @@ def __init__( if self.class_weights is not None: self.class_weights = self.class_weights.to(device).view(1, 1, -1) + small_topk = getattr(h, "stal_small_topk", 10) + small_topk_min = getattr(h, "stal_small_topk_min", 10) + # Global neutral defaults follow each branch's native matching budget. + # One-to-one supervision must not inherit one-to-many budget overrides. + if (small_topk, small_topk_min) == (10, 10) or tal_topk == 1 or tal_topk2 == 1: + small_topk = small_topk_min = tal_topk self.assigner = TaskAlignedAssigner( topk=tal_topk, num_classes=self.nc, @@ -375,10 +386,62 @@ def __init__( beta=6.0, stride=self.stride.tolist(), topk2=tal_topk2, + stal_enabled=getattr(h, "stal_enabled", False), + stal_stats=self.stal_stats, + stal_candidate_mode=getattr(h, "stal_candidate_mode", "fixed"), + stal_candidate_iou_floor=getattr(h, "stal_candidate_iou_floor", 0.0), + stal_expanded_quality_ratio=getattr(h, "stal_expanded_quality_ratio", 0.0), + stal_expanded_score_floor=getattr(h, "stal_expanded_score_floor", 1.0), + stal_area_threshold=getattr(h, "stal_area_threshold", 0.01), + stal_small_topk=small_topk, + stal_small_topk_min=small_topk_min, + stal_min_base_candidates=getattr(h, "stal_min_base_candidates", 0), + stal_max_extra_candidates=getattr(h, "stal_max_extra_candidates", 0), + stal_min_candidate_guarantee=getattr(h, "stal_min_candidate_guarantee", False), + stal_crowding_mode=getattr(h, "stal_crowding_mode", "none"), + stal_crowded_relaxation=getattr(h, "stal_crowded_relaxation", 0.0), + stal_relaxation=getattr(h, "stal_relaxation", 8.0), + stal_relaxation_scale_mode=getattr(h, "stal_relaxation_scale_mode", "constant"), + stal_relaxation_min=getattr(h, "stal_relaxation_min", 0.0), + stal_min_size_stride_ratio=getattr(h, "stal_min_size_stride_ratio", 0.0), + stal_warmup_epochs=getattr(h, "stal_warmup_epochs", 10.0), + stal_zero_positive_rescue=getattr(h, "stal_zero_positive_rescue", False), + stal_rescue_score_floor=getattr(h, "stal_rescue_score_floor", 0.0), + stal_rescue_floor_decay_epochs=getattr(h, "stal_rescue_floor_decay_epochs", 0.0), + stal_nwd_weight=getattr(h, "stal_nwd_weight", 0.0), + stal_nwd_constant=getattr(h, "stal_nwd_constant", 12.8), + stal_nwd_target_mode=getattr(h, "stal_nwd_target_mode", "match"), + stal_nwd_target_weight=getattr(h, "stal_nwd_target_weight", 0.0), + stal_nwd_zero_score_floor=getattr(h, "stal_nwd_zero_score_floor", 0.0), + stal_simd_weight=getattr(h, "stal_simd_weight", 0.0), ) self.bbox_loss = BboxLoss(m.reg_max).to(device) self.proj = torch.arange(m.reg_max, dtype=torch.float, device=device) + def pop_assignment_stats(self) -> torch.Tensor: + """Return and reset accumulated STAL/all-target assignment counts.""" + stats = self._assignment_stats.clone() + self._assignment_stats.zero_() + return stats + + def pop_rescue_stats(self) -> torch.Tensor: + """Return and reset accumulated zero-positive rescue diagnostics.""" + stats = self._rescue_stats.clone() + self._rescue_stats.zero_() + return stats + + def pop_assignment_stage_stats(self) -> torch.Tensor: + """Return and reset accumulated small-target assignment-stage diagnostics.""" + stats = self._assignment_stage_stats.clone() + self._assignment_stage_stats.zero_() + return stats + + def pop_target_score_stats(self) -> torch.Tensor: + """Return and reset target-score diagnostics for STAL, all, and absolute COCO-small foreground targets.""" + stats = self._target_score_stats.clone() + self._target_score_stats.zero_() + return stats + def preprocess(self, targets: torch.Tensor, batch_size: int, scale_tensor: torch.Tensor) -> torch.Tensor: """Preprocess targets by converting to tensor format and scaling coordinates.""" nl, ne = targets.shape @@ -437,7 +500,37 @@ def get_assigned_targets_and_loss(self, preds: dict[str, torch.Tensor], batch: d gt_labels, gt_bboxes, mask_gt, + image_size=imgsz, + epoch=batch.get("epoch", 0), ) + if self.stal_stats and "epoch" in batch: + self._assignment_stats += self.assigner.assignment_statistics( + gt_bboxes, mask_gt, fg_mask, target_gt_idx, image_size=imgsz + ) + self._rescue_stats += self.assigner.rescue_statistics().to(self.device) + self._assignment_stage_stats += self.assigner.assignment_stage_statistics().to(self.device) + foreground_scores = target_scores.sum(-1)[fg_mask] + foreground_gt = target_gt_idx[fg_mask] + foreground_batch = torch.arange(batch_size, device=self.device).unsqueeze(1).expand_as(fg_mask)[fg_mask] + assigned_boxes = gt_bboxes[foreground_batch, foreground_gt] + assigned_wh = (assigned_boxes[:, 2:] - assigned_boxes[:, :2]).clamp_min(0) + stal_foreground = self.assigner.small_target_mask(assigned_boxes, imgsz) + stal_scores = foreground_scores[stal_foreground] + coco_small = assigned_wh.prod(-1, dtype=torch.float32) < 32**2 + coco_small_scores = foreground_scores[coco_small] + self._target_score_stats += torch.stack( + ( + stal_foreground.sum(), + (stal_scores > 0).sum(), + stal_scores.double().sum(), + fg_mask.sum(), + (foreground_scores > 0).sum(), + foreground_scores.double().sum(), + coco_small.sum(), + (coco_small_scores > 0).sum(), + coco_small_scores.double().sum(), + ) + ) target_scores_sum = max(target_scores.sum(), 1) @@ -1212,6 +1305,28 @@ def __init__(self, model: torch.nn.Module, loss_fn=v8DetectionLoss): # final gain self.final_o2m = 0.1 + def _pop_assignment_diagnostic(self, name: str) -> torch.Tensor: + """Return one-to-many diagnostics and reset both assignment branches.""" + stats = getattr(self.one2many, name)() + getattr(self.one2one, name)() + return stats + + def pop_assignment_stats(self) -> torch.Tensor: + """Return and reset accumulated assignment statistics.""" + return self._pop_assignment_diagnostic("pop_assignment_stats") + + def pop_rescue_stats(self) -> torch.Tensor: + """Return and reset accumulated rescue statistics.""" + return self._pop_assignment_diagnostic("pop_rescue_stats") + + def pop_assignment_stage_stats(self) -> torch.Tensor: + """Return and reset accumulated assignment-stage statistics.""" + return self._pop_assignment_diagnostic("pop_assignment_stage_stats") + + def pop_target_score_stats(self) -> torch.Tensor: + """Return and reset accumulated target-score statistics.""" + return self._pop_assignment_diagnostic("pop_target_score_stats") + def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: """Calculate the sum of the loss for box, cls and dfl multiplied by batch size.""" preds = self.one2many.parse_output(preds) diff --git a/ultralytics/utils/tal.py b/ultralytics/utils/tal.py index 630d1d548..94c3d7879 100644 --- a/ultralytics/utils/tal.py +++ b/ultralytics/utils/tal.py @@ -2,8 +2,10 @@ from __future__ import annotations +import math + import torch -import torch.nn as nn +from torch import nn from . import LOGGER from .metrics import bbox_iou, probiou @@ -37,6 +39,34 @@ def __init__( stride: list | None = None, eps: float = 1e-9, topk2=None, + stal_enabled: bool = False, + stal_stats: bool = False, + stal_candidate_mode: str = "fixed", + stal_area_threshold: float = 0.01, + stal_small_topk: int | None = None, + stal_small_topk_min: int | None = None, + stal_min_base_candidates: int = 0, + stal_max_extra_candidates: int = 0, + stal_min_candidate_guarantee: bool = False, + stal_crowding_mode: str = "none", + stal_crowded_relaxation: float = 0.0, + stal_relaxation: float = 8.0, + stal_relaxation_scale_mode: str = "constant", + stal_relaxation_min: float = 0.0, + stal_min_size_stride_ratio: float = 0.0, + stal_warmup_epochs: float = 0.0, + stal_zero_positive_rescue: bool = False, + stal_rescue_score_floor: float = 0.0, + stal_rescue_floor_decay_epochs: float = 0.0, + stal_nwd_weight: float = 0.0, + stal_nwd_constant: float = 12.8, + stal_nwd_target_mode: str = "match", + stal_nwd_target_weight: float = 0.0, + stal_nwd_zero_score_floor: float = 0.0, + stal_candidate_iou_floor: float = 0.0, + stal_expanded_quality_ratio: float = 0.0, + stal_expanded_score_floor: float = 1.0, + stal_simd_weight: float = 0.0, ): """Initialize a TaskAlignedAssigner object with customizable hyperparameters. @@ -48,8 +78,76 @@ 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_enabled (bool, optional): Whether to enable small-target area-aware label assignment. + stal_stats (bool, optional): Whether to collect assignment diagnostics during training. + stal_candidate_mode (str, optional): Candidate policy: pure, fixed, or adaptive. + stal_area_threshold (float, optional): Maximum target-to-image area ratio that activates STAL. + stal_small_topk (int, optional): Adaptive-mode top-k for small targets; ``None`` follows ``topk``. + stal_small_topk_min (int, optional): Minimum top-k for near-zero-area targets. Values below the small-target + maximum enable square-root area interpolation. + stal_min_base_candidates (int, optional): Expand only small GTs with fewer base candidates. Zero disables. + stal_max_extra_candidates (int, optional): Preserve base TAL and add at most this many adaptive-only + candidates per small GT. Zero preserves the existing joint top-k behavior. + stal_min_candidate_guarantee (bool, optional): Give an otherwise candidate-free small GT its nearest + feature point while preserving normal TAL ranking and scoring. + stal_crowding_mode (str, optional): Reduce relaxation for small GTs whose full candidate masks overlap. + stal_crowded_relaxation (float, optional): Total width/height increase for crowded small GTs. + stal_relaxation (float, optional): Maximum total width and height added to the candidate region in pixels. + stal_relaxation_scale_mode (str, optional): Use a constant increase or interpolate it by square-root area. + stal_relaxation_min (float, optional): Total increase near the area threshold in square-root-area mode. + stal_min_size_stride_ratio (float, optional): Minimum candidate width and height as a multiple of the + model's minimum stride. Zero disables this guarantee. + stal_warmup_epochs (float, optional): Number of epochs used to linearly ramp STAL relaxation. + stal_zero_positive_rescue (bool, optional): Whether to rescue uncovered small GTs with their best legal + positive-quality candidate. + stal_rescue_score_floor (float, optional): Minimum supervision weight for a zero-quality bootstrap rescue. + stal_rescue_floor_decay_epochs (float, optional): Epochs used to linearly decay the bootstrap floor. + stal_nwd_weight (float, optional): NWD share in small-target assignment quality; zero preserves CIoU TAL. + stal_nwd_constant (float, optional): Positive normalization constant for NWD similarity. + stal_nwd_target_mode (str, optional): Use match quality or CIoU for final target-score normalization. + stal_nwd_target_weight (float, optional): NWD share for target scores in weighted mode. + stal_nwd_zero_score_floor (float, optional): Score floor for selected small anchors with zero target score. + stal_candidate_iou_floor (float, optional): Geometric IoU capacity floor for expanded small-GT candidates. + stal_expanded_quality_ratio (float, optional): Minimum alignment of an adaptive-only candidate relative to + the best candidate available before adaptive relaxation. Zero disables this gate. + stal_expanded_score_floor (float, optional): Minimum continuous target-score multiplier for adaptive-only + candidates relative to the best pre-relaxation candidate. One disables attenuation. + stal_simd_weight (float, optional): CIoU-to-SimD blend weight for small-target matching quality. """ super().__init__() + numeric_stal_parameters = { + "stal_area_threshold": stal_area_threshold, + "stal_small_topk": stal_small_topk, + "stal_small_topk_min": stal_small_topk_min, + "stal_min_base_candidates": stal_min_base_candidates, + "stal_max_extra_candidates": stal_max_extra_candidates, + "stal_crowded_relaxation": stal_crowded_relaxation, + "stal_relaxation": stal_relaxation, + "stal_relaxation_min": stal_relaxation_min, + "stal_min_size_stride_ratio": stal_min_size_stride_ratio, + "stal_warmup_epochs": stal_warmup_epochs, + "stal_rescue_score_floor": stal_rescue_score_floor, + "stal_rescue_floor_decay_epochs": stal_rescue_floor_decay_epochs, + "stal_nwd_weight": stal_nwd_weight, + "stal_nwd_constant": stal_nwd_constant, + "stal_nwd_target_weight": stal_nwd_target_weight, + "stal_nwd_zero_score_floor": stal_nwd_zero_score_floor, + "stal_candidate_iou_floor": stal_candidate_iou_floor, + "stal_expanded_quality_ratio": stal_expanded_quality_ratio, + "stal_expanded_score_floor": stal_expanded_score_floor, + "stal_simd_weight": stal_simd_weight, + } + for name, value in numeric_stal_parameters.items(): + if isinstance(value, bool): + raise TypeError(f"{name} must be numeric, not boolean") + for name, value in { + "stal_enabled": stal_enabled, + "stal_stats": stal_stats, + "stal_min_candidate_guarantee": stal_min_candidate_guarantee, + "stal_zero_positive_rescue": stal_zero_positive_rescue, + }.items(): + if not isinstance(value, bool): + raise TypeError(f"{name} must be a boolean, got {type(value).__name__}") self.topk = topk self.topk2 = topk2 or topk self.num_classes = num_classes @@ -58,9 +156,229 @@ def __init__( 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.eps = eps + if stal_min_candidate_guarantee and stal_candidate_iou_floor > 0: + raise ValueError("stal_min_candidate_guarantee cannot be combined with a positive stal_candidate_iou_floor") + for name, value in { + "stal_relaxation": stal_relaxation, + "stal_min_size_stride_ratio": stal_min_size_stride_ratio, + "stal_warmup_epochs": stal_warmup_epochs, + "stal_rescue_floor_decay_epochs": stal_rescue_floor_decay_epochs, + "stal_nwd_constant": stal_nwd_constant, + }.items(): + if not math.isfinite(value): + raise ValueError(f"{name} must be finite") + if not 0.0 <= stal_candidate_iou_floor <= 1.0: + raise ValueError("stal_candidate_iou_floor must be in [0, 1]") + if not 0.0 <= stal_expanded_quality_ratio <= 1.0: + raise ValueError("stal_expanded_quality_ratio must be in [0, 1]") + if not 0.0 <= stal_expanded_score_floor <= 1.0: + raise ValueError("stal_expanded_score_floor must be in [0, 1]") + if not 0.0 <= stal_simd_weight <= 1.0: + raise ValueError("stal_simd_weight must be in [0, 1]") + self.stal_candidate_iou_floor = float(stal_candidate_iou_floor) + self.stal_expanded_quality_ratio = float(stal_expanded_quality_ratio) + self.stal_expanded_score_floor = float(stal_expanded_score_floor) + self.stal_simd_weight = float(stal_simd_weight) + if not 0.0 < stal_area_threshold <= 1.0: + raise ValueError(f"stal_area_threshold must be in (0, 1], got {stal_area_threshold}") + if stal_small_topk is not None and ( + not isinstance(stal_small_topk, int) or isinstance(stal_small_topk, bool) or stal_small_topk < 1 + ): + raise ValueError(f"stal_small_topk must be a positive integer, got {stal_small_topk}") + effective_small_topk = topk if stal_small_topk is None else stal_small_topk + if stal_small_topk_min is not None and ( + not isinstance(stal_small_topk_min, int) + or isinstance(stal_small_topk_min, bool) + or not 1 <= stal_small_topk_min <= effective_small_topk + ): + raise ValueError(f"stal_small_topk_min must be in [1, stal_small_topk], got {stal_small_topk_min}") + if ( + not isinstance(stal_min_base_candidates, int) + or isinstance(stal_min_base_candidates, bool) + or stal_min_base_candidates < 0 + ): + raise ValueError(f"stal_min_base_candidates must be a non-negative integer, got {stal_min_base_candidates}") + if ( + not isinstance(stal_max_extra_candidates, int) + or isinstance(stal_max_extra_candidates, bool) + or stal_max_extra_candidates < 0 + ): + raise ValueError( + f"stal_max_extra_candidates must be a non-negative integer, got {stal_max_extra_candidates}" + ) + if stal_crowding_mode not in {"none", "candidate_overlap"}: + raise ValueError("stal_crowding_mode must be none or candidate_overlap") + if not 0.0 <= stal_crowded_relaxation <= stal_relaxation: + raise ValueError("stal_crowded_relaxation must be in [0, stal_relaxation]") + if stal_relaxation < 0.0: + raise ValueError(f"stal_relaxation must be non-negative, got {stal_relaxation}") + if stal_relaxation_scale_mode not in {"constant", "sqrt_area"}: + raise ValueError("stal_relaxation_scale_mode must be constant or sqrt_area") + if not 0.0 <= stal_relaxation_min <= stal_relaxation: + raise ValueError("stal_relaxation_min must be in [0, stal_relaxation]") + if stal_min_size_stride_ratio < 0.0: + raise ValueError("stal_min_size_stride_ratio must be non-negative") + if stal_min_size_stride_ratio > 0.0 and stal_relaxation > 0.0: + raise ValueError("stal_min_size_stride_ratio>0 requires stal_relaxation=0 for an isolated strategy") + if stal_min_size_stride_ratio > 0.0 and stal_relaxation_scale_mode != "constant": + raise ValueError("stal_min_size_stride_ratio>0 requires stal_relaxation_scale_mode=constant") + if stal_warmup_epochs < 0.0: + raise ValueError(f"stal_warmup_epochs must be non-negative, got {stal_warmup_epochs}") + if not 0.0 <= stal_rescue_score_floor <= 1.0: + raise ValueError(f"stal_rescue_score_floor must be in [0, 1], got {stal_rescue_score_floor}") + if stal_rescue_floor_decay_epochs < 0.0: + raise ValueError( + f"stal_rescue_floor_decay_epochs must be non-negative, got {stal_rescue_floor_decay_epochs}" + ) + if not 0.0 <= stal_nwd_weight <= 1.0: + raise ValueError(f"stal_nwd_weight must be in [0, 1], got {stal_nwd_weight}") + if stal_nwd_constant <= 0.0: + raise ValueError(f"stal_nwd_constant must be positive, got {stal_nwd_constant}") + if stal_nwd_target_mode not in {"match", "ciou", "weighted"}: + raise ValueError(f"stal_nwd_target_mode must be match, ciou, or weighted, got {stal_nwd_target_mode}") + if not 0.0 <= stal_nwd_target_weight <= stal_nwd_weight: + raise ValueError(f"stal_nwd_target_weight must be in [0, stal_nwd_weight], got {stal_nwd_target_weight}") + if not 0.0 <= stal_nwd_zero_score_floor <= 1.0: + raise ValueError(f"stal_nwd_zero_score_floor must be in [0, 1], got {stal_nwd_zero_score_floor}") + if stal_candidate_mode not in {"pure", "fixed", "adaptive"}: + raise ValueError(f"stal_candidate_mode must be pure, fixed, or adaptive, got {stal_candidate_mode}") + if stal_enabled and stal_candidate_mode == "pure": + raise ValueError("stal_enabled=True conflicts with stal_candidate_mode=pure") + self.stal_enabled = stal_enabled + self.stal_stats = stal_stats + self.stal_candidate_mode = "adaptive" if stal_enabled else stal_candidate_mode + if stal_rescue_score_floor > 0.0 and not stal_zero_positive_rescue: + raise ValueError("stal_rescue_score_floor>0 requires stal_zero_positive_rescue=True") + self.stal_area_threshold = float(stal_area_threshold) + self.stal_small_topk = effective_small_topk + self.stal_small_topk_min = effective_small_topk if stal_small_topk_min is None else stal_small_topk_min + self.stal_min_base_candidates = stal_min_base_candidates + self.stal_max_extra_candidates = stal_max_extra_candidates + self.stal_min_candidate_guarantee = stal_min_candidate_guarantee + self.stal_crowding_mode = stal_crowding_mode + self.stal_crowded_relaxation = float(stal_crowded_relaxation) + self.stal_relaxation = float(stal_relaxation) + self.stal_relaxation_scale_mode = stal_relaxation_scale_mode + self.stal_relaxation_min = float(stal_relaxation_min) + self.stal_min_size_stride_ratio = float(stal_min_size_stride_ratio) + self.stal_warmup_epochs = float(stal_warmup_epochs) + self.stal_zero_positive_rescue = stal_zero_positive_rescue + self.stal_rescue_score_floor = float(stal_rescue_score_floor) + self.stal_rescue_floor_decay_epochs = float(stal_rescue_floor_decay_epochs) + self.stal_nwd_weight = float(stal_nwd_weight) + self.stal_nwd_constant = float(stal_nwd_constant) + self.stal_nwd_target_mode = stal_nwd_target_mode + self.stal_nwd_target_weight = float(stal_nwd_target_weight) + self.stal_nwd_zero_score_floor = float(stal_nwd_zero_score_floor) + self._last_rescue_stats = torch.zeros(9, dtype=torch.long) + self._last_stage_stats = torch.zeros(9, dtype=torch.long) + self._last_rescue_targets = None + self._last_bootstrap_targets = None + self._last_bootstrap_mask = None + self._last_expanded_score_weights = None + + def rescue_statistics(self) -> torch.Tensor: + """Return diagnostics from the most recent zero-positive rescue attempt.""" + return self._last_rescue_stats.clone() + + def assignment_stage_statistics(self) -> torch.Tensor: + """Return small-target counts at candidate, alignment, top-k, and conflict stages.""" + return self._last_stage_stats.clone() + + def record_assignment_stages( + self, gt_bboxes, mask_gt, mask_in_gts, align_metric, pre_conflict_mask, post_conflict_mask, image_size + ) -> None: + """Record a behavior-neutral small-target assignment funnel for the latest batch.""" + image_size = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + valid = mask_gt.squeeze(-1).bool() + small = self.small_target_mask(gt_bboxes, image_size) & valid + has_legal = mask_in_gts.bool().any(-1) + has_nonzero_alignment = (align_metric > 0).any(-1) + has_positive_quality = (align_metric > self.eps).any(-1) + pre_selected = pre_conflict_mask.bool().any(-1) + post_selected = post_conflict_mask.bool().any(-1) + self._last_stage_stats = torch.stack( + ( + small.sum(), + (small & ~has_legal).sum(), + (small & has_legal & ~has_nonzero_alignment).sum(), + (small & has_nonzero_alignment & ~has_positive_quality).sum(), + (small & has_positive_quality).sum(), + (small & has_nonzero_alignment & ~pre_selected).sum(), + (small & pre_selected).sum(), + (small & pre_selected & ~post_selected).sum(), + (small & ~post_selected).sum(), + ) + ) + + def stal_rescue_score_floor_at_epoch(self, epoch: float) -> float: + """Return the bootstrap supervision floor after optional linear decay.""" + if self.stal_rescue_score_floor == 0.0 or self.stal_rescue_floor_decay_epochs == 0.0: + return self.stal_rescue_score_floor + remaining = 1.0 - min(max(float(epoch), 0.0) / self.stal_rescue_floor_decay_epochs, 1.0) + return self.stal_rescue_score_floor * remaining + + def stal_relaxation_at_epoch(self, epoch: float) -> float: + """Return the symmetric STAL size increase after applying linear warmup.""" + if self.stal_candidate_mode != "adaptive" or self.stal_relaxation == 0.0: + return 0.0 + if self.stal_warmup_epochs == 0.0: + return self.stal_relaxation + progress = min(max(float(epoch), 0.0) / self.stal_warmup_epochs, 1.0) + return self.stal_relaxation * progress + + def stal_min_candidate_size_at_epoch(self, epoch: float) -> float: + """Return the warmup-scaled minimum candidate width and height.""" + if self.stal_candidate_mode != "adaptive" or self.stal_min_size_stride_ratio == 0.0: + return 0.0 + progress = 1.0 if self.stal_warmup_epochs == 0.0 else min(max(float(epoch), 0.0) / self.stal_warmup_epochs, 1.0) + return self.stal_min_size_stride_ratio * min(self.stride) * progress + + def stal_relaxation_for_area( + self, relative_area: torch.Tensor, epoch: float, maximum: float | None = None + ) -> torch.Tensor: + """Return the warmup-scaled total size increase for each target area.""" + maximum = self.stal_relaxation_at_epoch(epoch) if maximum is None else float(maximum) + if self.stal_relaxation_scale_mode == "constant": + return torch.full_like(relative_area, maximum) + progress = maximum / self.stal_relaxation if self.stal_relaxation > 0.0 else 0.0 + minimum = self.stal_relaxation_min * progress + area_ratio = (relative_area / self.stal_area_threshold).clamp(0, 1).sqrt() + return minimum + (maximum - minimum) * (1 - area_ratio) + + def assignment_statistics(self, gt_bboxes, mask_gt, fg_mask, target_gt_idx, image_size): + """Return GT, positive, and zero-positive counts for STAL/COCO-size/all groups.""" + if gt_bboxes.shape[1] == 0: + return torch.zeros(15, dtype=torch.long, device=gt_bboxes.device) + image_size = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + if image_size.numel() != 2 or (image_size <= 0).any(): + raise ValueError(f"image_size must contain positive height and width, got {image_size.tolist()}") + gt_wh = (gt_bboxes[..., 2:] - gt_bboxes[..., :2]).clamp_min_(0) + valid_gt = mask_gt.squeeze(-1).bool() + area = gt_wh.prod(-1, dtype=torch.float32) + stal_gt = (area / image_size.prod(dtype=torch.float32) < self.stal_area_threshold) & valid_gt + foreground = fg_mask.bool() + positives_per_gt = torch.zeros_like(valid_gt, dtype=torch.long) + positives_per_gt.scatter_add_(1, target_gt_idx.long(), foreground.long()) + + def group_counts(group): + return torch.stack((group.sum(), positives_per_gt[group].sum(), (group & (positives_per_gt == 0)).sum())) + + small = (area < 32**2) & valid_gt + medium = (area >= 32**2) & (area < 96**2) & valid_gt + large = (area >= 96**2) & valid_gt + return torch.cat( + ( + group_counts(stal_gt), + group_counts(small), + group_counts(medium), + group_counts(large), + group_counts(valid_gt), + ) + ) @torch.no_grad() - def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt): + def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt, image_size=None, epoch=0): """Compute the task-aligned assignment. Args: @@ -70,6 +388,8 @@ def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_g gt_labels (torch.Tensor): Ground truth labels with shape (bs, n_max_boxes, 1). gt_bboxes (torch.Tensor): Ground truth boxes with shape (bs, n_max_boxes, 4). mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (bs, n_max_boxes, 1). + image_size (tuple | torch.Tensor, optional): Current input image height and width. + epoch (int, optional): Zero-based training epoch used by STAL warmup. Returns: target_labels (torch.Tensor): Target labels with shape (bs, num_total_anchors). @@ -84,28 +404,40 @@ def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_g self.bs = pd_scores.shape[0] self.n_max_boxes = gt_bboxes.shape[1] device = gt_bboxes.device + self._last_rescue_stats = torch.zeros(9, dtype=torch.long, device=device) + self._last_stage_stats = torch.zeros(9, dtype=torch.long, device=device) + self._last_rescue_targets = None + self._last_bootstrap_targets = None + self._last_bootstrap_mask = None if self.n_max_boxes == 0: return ( - torch.full_like(pd_scores[..., 0], self.num_classes), + torch.full_like(pd_scores[..., 0], self.num_classes, dtype=torch.long), torch.zeros_like(pd_bboxes), torch.zeros_like(pd_scores), - torch.zeros_like(pd_scores[..., 0]), - torch.zeros_like(pd_scores[..., 0]), + torch.zeros_like(pd_scores[..., 0], dtype=torch.bool), + torch.zeros_like(pd_scores[..., 0], dtype=torch.long), ) try: - return self._forward(pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt) + return self._forward( + pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt, image_size=image_size, epoch=epoch + ) except RuntimeError as e: if "out of memory" not in str(e).lower(): raise # Recover outside the except block: exiting it drops e.__traceback__, releasing the failed attempt's GPU # intermediates back to the allocator so the copy-back below can succeed LOGGER.warning("CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU") - result = self._forward(*(t.cpu() for t in (pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt))) + cpu_image_size = image_size.cpu() if isinstance(image_size, torch.Tensor) else image_size + result = self._forward( + *(t.cpu() for t in (pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt)), + image_size=cpu_image_size, + epoch=epoch, + ) return tuple(t.to(device) for t in result) - def _forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt): + def _forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt, image_size=None, epoch=0): """Compute the task-aligned assignment. Args: @@ -123,13 +455,43 @@ def _forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_ fg_mask (torch.Tensor): Foreground mask with shape (bs, num_total_anchors). target_gt_idx (torch.Tensor): Target ground truth indices with shape (bs, num_total_anchors). """ - mask_pos, align_metric, overlaps = self.get_pos_mask( - pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt + mask_pos, align_metric, overlaps, mask_in_gts, target_overlaps = self.get_pos_mask( + pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt, image_size=image_size, epoch=epoch ) + pre_conflict_mask = mask_pos target_gt_idx, fg_mask, mask_pos = self.select_highest_overlaps( mask_pos, overlaps, self.n_max_boxes, align_metric ) + if self.stal_stats: + self.record_assignment_stages( + gt_bboxes, mask_gt, mask_in_gts, align_metric, pre_conflict_mask, mask_pos, image_size + ) + if self.stal_zero_positive_rescue: + mask_pos = self.rescue_zero_positive_small_targets( + mask_pos, + align_metric, + mask_in_gts, + gt_bboxes, + mask_gt, + image_size, + anc_points, + epoch, + available_anchors=~fg_mask.bool(), + ) + target_gt_idx, fg_mask, mask_pos = self.select_highest_overlaps( + mask_pos, overlaps, self.n_max_boxes, align_metric + ) + rescued = (self._last_rescue_targets & (mask_pos.sum(-1) > 0)).sum() + self._last_rescue_stats[5] = rescued + self._last_rescue_stats[6] = self._last_rescue_stats[4] - rescued + bootstrap_rescued = (self._last_bootstrap_targets & (mask_pos.sum(-1) > 0)).sum() + self._last_rescue_stats[8] = bootstrap_rescued + self._last_bootstrap_mask &= mask_pos.bool() + + expanded_score_weights = None + if self._last_expanded_score_weights is not None: + expanded_score_weights = (self._last_expanded_score_weights * mask_pos).amax(-2).unsqueeze(-1) # Assigned target target_labels, target_bboxes, target_scores = self.get_targets(gt_labels, gt_bboxes, target_gt_idx, fg_mask) @@ -137,13 +499,31 @@ def _forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_ # Normalize align_metric *= mask_pos pos_align_metrics = align_metric.amax(dim=-1, keepdim=True) # b, max_num_obj - pos_overlaps = (overlaps * mask_pos).amax(dim=-1, keepdim=True) # b, max_num_obj + pos_overlaps = (target_overlaps * mask_pos).amax(dim=-1, keepdim=True) # b, max_num_obj norm_align_metric = (align_metric * pos_overlaps / (pos_align_metrics + self.eps)).amax(-2).unsqueeze(-1) + if self.stal_nwd_zero_score_floor > 0.0: + assigned_boxes = gt_bboxes.gather(1, target_gt_idx.unsqueeze(-1).expand(-1, -1, 4)) + image_size_tensor = torch.as_tensor(image_size, dtype=assigned_boxes.dtype, device=assigned_boxes.device) + assigned_small = self.small_target_mask(assigned_boxes, image_size_tensor) + zero_score_small = fg_mask.bool() & assigned_small & (norm_align_metric.squeeze(-1) <= self.eps) + norm_align_metric = torch.where( + zero_score_small.unsqueeze(-1), + norm_align_metric.new_tensor(self.stal_nwd_zero_score_floor), + norm_align_metric, + ) + if self._last_bootstrap_mask is not None: + bootstrap_anchor_mask = self._last_bootstrap_mask.any(dim=1).unsqueeze(-1) + floor = self.stal_rescue_score_floor_at_epoch(epoch) + norm_align_metric = torch.where( + bootstrap_anchor_mask, norm_align_metric.clamp_min(floor), norm_align_metric + ) + if expanded_score_weights is not None: + norm_align_metric = norm_align_metric * expanded_score_weights target_scores = target_scores * norm_align_metric return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx - def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt): + def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt, image_size=None, epoch=0): """Get positive mask for each ground truth box. Args: @@ -158,18 +538,172 @@ def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, m mask_pos (torch.Tensor): Positive mask with shape (bs, max_num_obj, h*w). align_metric (torch.Tensor): Alignment metric with shape (bs, max_num_obj, h*w). overlaps (torch.Tensor): Overlaps between predicted vs ground truth boxes with shape (bs, max_num_obj, h*w). + mask_in_gts (torch.Tensor): Legal candidate-region mask with shape (bs, max_num_obj, h*w). """ - mask_in_gts = self.select_candidates_in_gts(anc_points, gt_bboxes, mask_gt) + mask_in_gts = self.select_candidates_in_gts(anc_points, gt_bboxes, mask_gt, image_size=image_size, epoch=epoch) # 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()) + align_metric, overlaps, target_overlaps = self.get_box_metrics( + pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_in_gts * mask_gt, image_size=image_size + ) + self._last_expanded_score_weights = None + if self.stal_candidate_mode == "adaptive" and self.stal_expanded_score_floor < 1.0: + if image_size is None: + raise ValueError("image_size is required when adaptive expanded-candidate score weighting is active") + base_candidates = self.select_candidates_in_gts( + anc_points, gt_bboxes, mask_gt, image_size=image_size, epoch=epoch, relaxation_override=0.0 + ) + valid_gt = mask_gt.squeeze(-1).bool() + image_size_tensor = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + small_gt = self.small_target_mask(gt_bboxes, image_size_tensor) & valid_gt + has_base_candidate = base_candidates.any(-1) + best_base_metric = align_metric.masked_fill(~base_candidates, -torch.inf).amax(-1) + quality_ratio = (align_metric / best_base_metric.clamp_min(self.eps).unsqueeze(-1)).clamp(0, 1) + soft_weight = self.stal_expanded_score_floor + (1.0 - self.stal_expanded_score_floor) * quality_ratio + attenuated = small_gt.unsqueeze(-1) & has_base_candidate.unsqueeze(-1) & ~base_candidates + self._last_expanded_score_weights = torch.where(attenuated, soft_weight, torch.ones_like(soft_weight)) + if self.stal_candidate_mode == "adaptive" and self.stal_expanded_quality_ratio > 0.0: + if image_size is None: + raise ValueError("image_size is required when adaptive expanded-candidate quality gating is active") + base_candidates = self.select_candidates_in_gts( + anc_points, gt_bboxes, mask_gt, image_size=image_size, epoch=epoch, relaxation_override=0.0 + ) + valid_gt = mask_gt.squeeze(-1).bool() + image_size_tensor = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + small_gt = self.small_target_mask(gt_bboxes, image_size_tensor) & valid_gt + has_base_candidate = base_candidates.any(-1) + best_base_metric = align_metric.masked_fill(~base_candidates, -torch.inf).amax(-1) + # A relative threshold is undefined when every base candidate has zero alignment. Without the + # positive-base guard, ``0 >= 0 * ratio`` admits every zero-quality expanded candidate. + has_positive_base = best_base_metric > self.eps + competitive = has_positive_base.unsqueeze(-1) & ( + align_metric >= best_base_metric.unsqueeze(-1) * self.stal_expanded_quality_ratio + ) + gated_gt = (small_gt & has_base_candidate).unsqueeze(-1) + quality_candidates = base_candidates | competitive + mask_in_gts = torch.where(gated_gt, mask_in_gts & quality_candidates, mask_in_gts) + align_metric = align_metric * mask_in_gts + # Keep standard TAL top-k for non-small GTs; optionally interpolate small-target top-k from relative area. + valid_gt = mask_gt.squeeze(-1).bool() + base_topk_mask = valid_gt.unsqueeze(-1).expand(-1, -1, self.topk) + mask_topk = self.select_topk_candidates(align_metric, topk_mask=base_topk_mask) + if self.stal_candidate_mode == "adaptive" and self.stal_max_extra_candidates > 0: + if image_size is None: + raise ValueError("image_size is required when adaptive extra-candidate limiting is active") + base_candidates = self.select_candidates_in_gts( + anc_points, gt_bboxes, mask_gt, image_size=image_size, epoch=epoch, relaxation_override=0.0 + ) + image_size_tensor = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + small_gt = self.small_target_mask(gt_bboxes, image_size_tensor) & valid_gt + base_metrics = align_metric.masked_fill(~base_candidates, -torch.inf) + base_selected = self.select_topk_candidates(base_metrics, topk_mask=base_topk_mask) * base_candidates + extra_candidates = mask_in_gts.bool() & ~base_candidates + extra_slots = torch.ones( + (*valid_gt.shape, self.stal_max_extra_candidates), dtype=torch.bool, device=valid_gt.device + ) + extra_slots &= valid_gt.unsqueeze(-1) + extra_metrics = align_metric.masked_fill(~extra_candidates, -torch.inf) + extra_selected = self.select_topk_candidates(extra_metrics, topk_mask=extra_slots) * extra_candidates + limited_selection = torch.maximum(base_selected, extra_selected) + mask_topk = torch.where(small_gt.unsqueeze(-1), limited_selection, mask_topk) + adaptive_topk_active = self.stal_small_topk != self.topk or self.stal_small_topk_min != self.topk + if self.stal_candidate_mode == "adaptive" and adaptive_topk_active: + if image_size is None: + raise ValueError("image_size is required when adaptive small-target top-k is active") + image_size_tensor = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + small_gt = self.small_target_mask(gt_bboxes, image_size_tensor) & valid_gt + gt_area = (gt_bboxes[..., 2:] - gt_bboxes[..., :2]).clamp_min(0).prod(-1, dtype=torch.float32) + area_fraction = ( + (gt_area / image_size_tensor.prod(dtype=torch.float32) / self.stal_area_threshold).clamp(0, 1).sqrt() + ) + dynamic_topk = ( + (self.stal_small_topk_min + (self.stal_small_topk - self.stal_small_topk_min) * area_fraction) + .round() + .long() + ) + topk_slots = torch.arange(self.stal_small_topk, device=gt_bboxes.device) + small_topk_mask = valid_gt.unsqueeze(-1) & (topk_slots < dynamic_topk.unsqueeze(-1)) + adaptive_topk = self.select_topk_candidates(align_metric, topk_mask=small_topk_mask) + if self.stal_max_extra_candidates > 0: + # Both policies constrain the result; adaptive top-k cannot undo the extra-candidate cap. + adaptive_topk = adaptive_topk * limited_selection + mask_topk = torch.where(small_gt.unsqueeze(-1), adaptive_topk, mask_topk) # Merge all mask to a final mask, (b, max_num_obj, h*w) mask_pos = mask_topk * mask_in_gts * mask_gt - return mask_pos, align_metric, overlaps + return mask_pos, align_metric, overlaps, mask_in_gts, target_overlaps + + def rescue_zero_positive_small_targets( + self, + mask_pos, + align_metric, + mask_in_gts, + gt_bboxes, + mask_gt, + image_size, + anc_points=None, + epoch=0, + available_anchors=None, + ): + """Propose a quality-ranked or center-prior bootstrap candidate for each uncovered small GT.""" + if not self.stal_zero_positive_rescue: + return mask_pos + if image_size is None: + raise ValueError("image_size is required when stal_zero_positive_rescue=True") + + image_size = torch.as_tensor(image_size, dtype=gt_bboxes.dtype, device=gt_bboxes.device) + if image_size.numel() != 2 or (image_size <= 0).any(): + raise ValueError(f"image_size must contain positive height and width, got {image_size.tolist()}") + + small_gt = self.small_target_mask(gt_bboxes, image_size) & mask_gt.squeeze(-1).bool() + uncovered_small_gt = small_gt & (mask_pos.sum(-1) == 0) + legal_candidates = mask_in_gts.bool() & mask_gt.bool() + has_legal = legal_candidates.any(-1) + if available_anchors is not None: + legal_candidates &= available_anchors.bool().unsqueeze(1) + has_free_legal = legal_candidates.any(-1) + legal_metrics = align_metric.masked_fill(~legal_candidates, -torch.inf) + best_metric, quality_best_idx = legal_metrics.max(-1, keepdim=True) + positive_quality = torch.isfinite(best_metric.squeeze(-1)) & (best_metric.squeeze(-1) > self.eps) + floor = self.stal_rescue_score_floor_at_epoch(epoch) + bootstrap_gt = uncovered_small_gt & has_free_legal & ~positive_quality & (floor > 0.0) + if bootstrap_gt.any(): + if anc_points is None: + raise ValueError("anc_points is required when stal_rescue_score_floor>0") + gt_centers = (gt_bboxes[..., :2] + gt_bboxes[..., 2:]) / 2 + gt_wh = (gt_bboxes[..., 2:] - gt_bboxes[..., :2]).clamp_min(0) + scale = gt_wh.clamp_min(1.0).unsqueeze(-2) + distance = ((anc_points.view(1, 1, -1, 2) - gt_centers.unsqueeze(-2)) / scale).square().sum(-1) + bootstrap_best_idx = distance.masked_fill(~legal_candidates, torch.inf).argmin(-1, keepdim=True) + else: + bootstrap_best_idx = quality_best_idx + best_idx = torch.where(positive_quality.unsqueeze(-1), quality_best_idx, bootstrap_best_idx) + eligible_gt = uncovered_small_gt & (positive_quality | bootstrap_gt) + eligible = eligible_gt.unsqueeze(-1) + + self._last_rescue_targets = uncovered_small_gt + self._last_bootstrap_targets = bootstrap_gt + self._last_rescue_stats = torch.stack( + ( + uncovered_small_gt.sum(), + (uncovered_small_gt & has_legal).sum(), + (uncovered_small_gt & has_free_legal).sum(), + (uncovered_small_gt & has_free_legal & positive_quality).sum(), + eligible_gt.sum(), + torch.zeros((), dtype=torch.long, device=mask_pos.device), + torch.zeros((), dtype=torch.long, device=mask_pos.device), + bootstrap_gt.sum(), + torch.zeros((), dtype=torch.long, device=mask_pos.device), + ) + ) + + rescue = torch.zeros_like(mask_pos) + rescue.scatter_(-1, best_idx, eligible.to(mask_pos.dtype)) + bootstrap_rescue = torch.zeros_like(mask_pos, dtype=torch.bool) + bootstrap_rescue.scatter_(-1, best_idx, bootstrap_gt.unsqueeze(-1)) + self._last_bootstrap_mask = bootstrap_rescue + return torch.maximum(mask_pos, rescue) - def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt): + def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt, image_size=None): """Compute alignment metric given predicted and ground truth bounding boxes. Args: @@ -181,11 +715,13 @@ def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt): Returns: align_metric (torch.Tensor): Alignment metric combining classification and localization. - overlaps (torch.Tensor): IoU overlaps between predicted and ground truth boxes. + overlaps (torch.Tensor): Matching quality, using CIoU or a size-gated CIoU/NWD blend. + target_overlaps (torch.Tensor): Quality used to normalize target scores. """ na = pd_bboxes.shape[-2] mask_gt = mask_gt.bool() # b, max_num_obj, h*w overlaps = torch.zeros([self.bs, self.n_max_boxes, na], dtype=pd_bboxes.dtype, device=pd_bboxes.device) + target_overlaps = torch.zeros_like(overlaps) bbox_scores = torch.zeros([self.bs, self.n_max_boxes, na], dtype=pd_scores.dtype, device=pd_scores.device) # Do not boolean-index expanded views here. On MPS, the backend can return different numbers of @@ -198,10 +734,68 @@ def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt): bbox_scores[batch_idx, gt_idx, anchor_idx] = pd_scores[batch_idx, anchor_idx, labels] pd_boxes = pd_bboxes[batch_idx, anchor_idx] gt_boxes = gt_bboxes[batch_idx, gt_idx] - overlaps[batch_idx, gt_idx, anchor_idx] = self.iou_calculation(gt_boxes, pd_boxes) + ciou = self.iou_calculation(gt_boxes, pd_boxes) + quality = ciou + small_gt = None + if self.stal_nwd_weight > 0.0: + if image_size is None: + raise ValueError("image_size is required when stal_nwd_weight>0") + image_size_tensor = torch.as_tensor(image_size, dtype=gt_boxes.dtype, device=gt_boxes.device) + if image_size_tensor.numel() != 2 or (image_size_tensor <= 0).any(): + raise ValueError( + f"image_size must contain positive height and width, got {image_size_tensor.tolist()}" + ) + small_gt = self.small_target_mask(gt_boxes, image_size_tensor) + nwd = self.nwd_similarity(gt_boxes, pd_boxes) + blended = quality.lerp(nwd, self.stal_nwd_weight) + quality = torch.where(small_gt, blended, quality) + if self.stal_simd_weight > 0.0: + if image_size is None: + raise ValueError("image_size is required when stal_simd_weight>0") + if small_gt is None: + image_size_tensor = torch.as_tensor(image_size, dtype=gt_boxes.dtype, device=gt_boxes.device) + if image_size_tensor.numel() != 2 or (image_size_tensor <= 0).any(): + raise ValueError("image_size must contain positive height and width") + small_gt = self.small_target_mask(gt_boxes, image_size_tensor) + simd = self.simd_similarity(gt_boxes, pd_boxes) + quality = torch.where(small_gt, quality.lerp(simd, self.stal_simd_weight), quality) + overlaps[batch_idx, gt_idx, anchor_idx] = quality + if self.stal_nwd_target_mode == "match": + target_quality = quality + elif self.stal_nwd_target_mode == "weighted" and self.stal_nwd_target_weight > 0.0: + target_blend = ciou.lerp(nwd, self.stal_nwd_target_weight) + target_quality = torch.where(small_gt, target_blend, ciou) + else: + target_quality = ciou + target_overlaps[batch_idx, gt_idx, anchor_idx] = target_quality align_metric = bbox_scores.pow(self.alpha) * overlaps.pow(self.beta) - return align_metric, overlaps + return align_metric, overlaps, target_overlaps + + def nwd_similarity(self, gt_bboxes, pd_bboxes): + """Return normalized Gaussian Wasserstein similarity for paired horizontal xyxy boxes.""" + gt_xywh = xyxy2xywh(gt_bboxes) + pd_xywh = xyxy2xywh(pd_bboxes) + wasserstein_2 = (gt_xywh[..., :2] - pd_xywh[..., :2]).square().sum(-1) + wasserstein_2 += ((gt_xywh[..., 2:] - pd_xywh[..., 2:]) / 2).square().sum(-1) + return torch.exp(-torch.sqrt(wasserstein_2.clamp_min(0)) / self.stal_nwd_constant) + + @staticmethod + def simd_similarity(gt_bboxes, pd_bboxes): + """Return paired SimD using the official VisDrone implementation's x=6.13 and y=4.59 constants.""" + gt = xyxy2xywh(gt_bboxes) + pred = xyxy2xywh(pd_bboxes) + sums = (gt[..., 2:] + pred[..., 2:]).clamp_min(1e-9) + location_scale = sums / sums.new_tensor((6.13, 4.59)) + sim_location = ((gt[..., :2] - pred[..., :2]) / location_scale).square().sum(-1).sqrt() + sim_shape = ((gt[..., 2:] - pred[..., 2:]) / location_scale).square().sum(-1).sqrt() + return torch.exp(-(sim_location + sim_shape)) + + def small_target_mask(self, gt_bboxes, image_size): + """Return the existing strict relative-area STAL gate for paired horizontal xyxy boxes.""" + gt_wh = (gt_bboxes[..., 2:] - gt_bboxes[..., :2]).clamp_min(0) + # Area products are evaluated in FP32 because typical 800x800 images overflow FP16 before division. + return gt_wh.prod(-1, dtype=torch.float32) / image_size.prod(dtype=torch.float32) < self.stal_area_threshold def iou_calculation(self, gt_bboxes, pd_bboxes): """Calculate IoU for horizontal bounding boxes. @@ -229,18 +823,20 @@ def select_topk_candidates(self, metrics, topk_mask=None): (torch.Tensor): A tensor of shape (b, max_num_obj, h*w) containing the selected top-k candidates. """ # (b, max_num_obj, topk) - topk_metrics, topk_idxs = torch.topk(metrics, self.topk, dim=-1, largest=True) + candidate_topk = self.topk if topk_mask is None else topk_mask.shape[-1] + candidate_topk = min(candidate_topk, metrics.shape[-1]) + if topk_mask is not None: + topk_mask = topk_mask[..., :candidate_topk] + topk_metrics, topk_idxs = torch.topk(metrics, candidate_topk, dim=-1, largest=True) if topk_mask is None: topk_mask = (topk_metrics.max(-1, keepdim=True)[0] > self.eps).expand_as(topk_idxs) # (b, max_num_obj, topk) - topk_idxs.masked_fill_(~topk_mask, 0) - # (b, max_num_obj, topk, h*w) -> (b, max_num_obj, h*w) count_tensor = torch.zeros(metrics.shape, dtype=torch.int8, device=topk_idxs.device) - ones = torch.ones_like(topk_idxs[:, :, :1], dtype=torch.int8, device=topk_idxs.device) - for k in range(self.topk): - # Expand topk_idxs for each value of k and add 1 at the specified positions - count_tensor.scatter_add_(-1, topk_idxs[:, :, k : k + 1], ones) + for k in range(candidate_topk): + # Disabled variable-top-k slots add zero instead of aliasing a real candidate at anchor index zero. + slot_enabled = topk_mask[:, :, k : k + 1].to(torch.int8) + count_tensor.scatter_add_(-1, topk_idxs[:, :, k : k + 1], slot_enabled) # Filter invalid bboxes count_tensor.masked_fill_(count_tensor > 1, 0) @@ -286,7 +882,9 @@ def get_targets(self, gt_labels, gt_bboxes, target_gt_idx, fg_mask): return target_labels, target_bboxes, target_scores - def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt, eps=1e-9): + def select_candidates_in_gts( + self, xy_centers, gt_bboxes, mask_gt, eps=1e-9, image_size=None, epoch=0, relaxation_override=None + ): """Select positive anchor centers within ground truth bounding boxes. Args: @@ -294,6 +892,9 @@ def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt, eps=1e-9): gt_bboxes (torch.Tensor): Ground truth bounding boxes, shape (b, n_boxes, 4). mask_gt (torch.Tensor): Mask for valid ground truth boxes, shape (b, n_boxes, 1). eps (float, optional): Small value for numerical stability. + image_size (tuple | torch.Tensor, optional): Current input image height and width. + epoch (int, optional): Zero-based training epoch used by STAL warmup. + relaxation_override (float, optional): Internal override used to recover the pre-relaxation candidate set. Returns: (torch.Tensor): Boolean mask of positive anchors, shape (b, n_boxes, h*w). @@ -303,16 +904,97 @@ def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt, eps=1e-9): - 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) - - 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) + relaxation = self.stal_relaxation_at_epoch(epoch) if relaxation_override is None else relaxation_override + min_candidate_size = self.stal_min_candidate_size_at_epoch(epoch) if relaxation_override is None else 0.0 + stal_mask = None + if self.stal_candidate_mode == "adaptive" and (relaxation > 0.0 or min_candidate_size > 0.0): + if image_size is None: + raise ValueError("image_size is required when adaptive STAL candidate geometry is active") + image_size = torch.as_tensor(image_size, dtype=gt_bboxes_xywh.dtype, device=gt_bboxes_xywh.device) + if image_size.numel() != 2 or (image_size <= 0).any(): + raise ValueError(f"image_size must contain positive height and width, got {image_size.tolist()}") + relative_area = gt_bboxes_xywh[..., 2:].prod(-1, keepdim=True, dtype=torch.float32) / image_size.prod( + dtype=torch.float32 + ) + stal_mask = (relative_area < self.stal_area_threshold) & mask_gt.bool() + + if self.stal_candidate_mode != "pure": + 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:], + ) + if stal_mask is not None and self.stal_min_base_candidates > 0: + base_boxes = xywh2xyxy(gt_bboxes_xywh) + base_lt, base_rb = base_boxes.unsqueeze(2).chunk(2, 3) + base_candidates = ((xy_centers - base_lt > eps) & (base_rb - xy_centers > eps)).all(3) + undercovered = base_candidates.sum(-1, keepdim=True) < self.stal_min_base_candidates + stal_mask &= undercovered + if stal_mask is not None: + if min_candidate_size > 0.0: + min_wh = torch.full_like(gt_bboxes_xywh[..., 2:], min_candidate_size) + gt_bboxes_xywh[..., 2:] = torch.where( + stal_mask.expand_as(gt_bboxes_xywh[..., 2:]), + torch.maximum(gt_bboxes_xywh[..., 2:], min_wh), + gt_bboxes_xywh[..., 2:], + ) + relaxation_per_gt = self.stal_relaxation_for_area(relative_area, epoch, maximum=relaxation) + if relaxation > 0.0 and self.stal_crowding_mode == "candidate_overlap": + full_wh = torch.where( + stal_mask.expand_as(gt_bboxes_xywh[..., 2:]), + gt_bboxes_xywh[..., 2:] + relaxation, + gt_bboxes_xywh[..., 2:], + ) + full_boxes = xywh2xyxy(torch.cat((gt_bboxes_xywh[..., :2], full_wh), dim=-1)) + full_lt, full_rb = full_boxes.unsqueeze(2).chunk(2, 3) + full_candidates = ((xy_centers - full_lt > eps) & (full_rb - xy_centers > eps)).all(3) + full_candidates &= mask_gt.bool() + shared = full_candidates & (full_candidates.sum(1, keepdim=True) > 1) + crowded = shared.any(-1, keepdim=True) & stal_mask + crowded_relaxation = self.stal_crowded_relaxation * relaxation / self.stal_relaxation + relaxation_per_gt = torch.where( + crowded, relaxation_per_gt.clamp(max=crowded_relaxation), relaxation_per_gt + ) + if relaxation > 0.0: + gt_bboxes_xywh[..., 2:] = torch.where( + stal_mask.expand_as(gt_bboxes_xywh[..., 2:]), + gt_bboxes_xywh[..., 2:] + relaxation_per_gt, + gt_bboxes_xywh[..., 2:], + ) + expanded_boxes = xywh2xyxy(gt_bboxes_xywh) + + lt, rb = expanded_boxes.unsqueeze(2).chunk(2, 3) # (b, n_boxes, 1, 2) left-top, right-bottom + candidates = ((xy_centers - lt > eps) & (rb - xy_centers > eps)).all(3) + if self.stal_candidate_iou_floor > 0.0 and self.stal_candidate_mode != "pure": + if image_size is None: + raise ValueError("image_size is required when stal_candidate_iou_floor>0") + image_size = torch.as_tensor(image_size, device=gt_bboxes.device, dtype=gt_bboxes.dtype) + if image_size.numel() != 2 or (image_size <= 0).any(): + raise ValueError("image_size must contain positive height and width") + small = self.small_target_mask(gt_bboxes, image_size) & mask_gt.squeeze(-1).bool() + # Nonnegative DFL distances force a predicted box to contain its feature point. The smallest + # rectangle enclosing both that point and the original GT gives an optimistic IoU ceiling. + lower = torch.minimum(gt_bboxes[..., None, :2], xy_centers) + upper = torch.maximum(gt_bboxes[..., None, 2:], xy_centers) + gt_area = (gt_bboxes[..., 2:] - gt_bboxes[..., :2]).clamp_min(0).prod(-1, dtype=torch.float32) + enclosing_area = (upper - lower).prod(-1, dtype=torch.float32).clamp_min(self.eps) + capacity = gt_area.unsqueeze(-1) / enclosing_area + candidates &= ~small.unsqueeze(-1) | (capacity >= self.stal_candidate_iou_floor) + if self.stal_min_candidate_guarantee and self.stal_candidate_mode == "adaptive": + if image_size is None: + raise ValueError("image_size is required when stal_min_candidate_guarantee=True") + image_size = torch.as_tensor(image_size, device=gt_bboxes.device, dtype=gt_bboxes.dtype) + small = self.small_target_mask(gt_bboxes, image_size) & mask_gt.squeeze(-1).bool() + candidate_free = small & ~candidates.any(-1) + if candidate_free.any(): + gt_centers = (gt_bboxes[..., :2] + gt_bboxes[..., 2:]) / 2 + distance = (xy_centers.view(1, 1, -1, 2) - gt_centers.unsqueeze(-2)).square().sum(-1) + nearest = distance.argmin(-1, keepdim=True) + guarantee = torch.zeros_like(candidates) + guarantee.scatter_(-1, nearest, candidate_free.unsqueeze(-1)) + candidates |= guarantee + return candidates 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. @@ -333,6 +1015,8 @@ def select_highest_overlaps(self, mask_pos, overlaps, n_max_boxes, align_metric) if fg_mask.max() > 1: # one anchor is assigned to multiple gt_bboxes mask_multi_gts = (fg_mask.unsqueeze(1) > 1).expand(-1, n_max_boxes, -1) # (b, n_max_boxes, h*w) + # Resolve conflicts only among GTs that actually nominated the anchor. Taking argmax over every GT can + # reassign an anchor to a non-candidate GT whose IoU happens to be larger, breaking the top-k contract. candidate_overlaps = overlaps.masked_fill(~mask_pos.bool(), -torch.inf) max_overlaps_idx = candidate_overlaps.argmax(1) # (b, h*w) is_max_overlaps = torch.zeros(mask_pos.shape, dtype=mask_pos.dtype, device=mask_pos.device) @@ -341,8 +1025,8 @@ def select_highest_overlaps(self, mask_pos, overlaps, n_max_boxes, align_metric) fg_mask = mask_pos.sum(-2) - if self.topk2 != self.topk: - align_metric = align_metric * mask_pos # update overlaps + if self.topk2 != self.topk or self.topk2 == 1: + 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 topk_idx = torch.zeros(mask_pos.shape, dtype=mask_pos.dtype, device=mask_pos.device) # update mask_pos @@ -361,7 +1045,7 @@ def iou_calculation(self, gt_bboxes, pd_bboxes): """Calculate IoU for rotated bounding boxes.""" return probiou(gt_bboxes, pd_bboxes).squeeze(-1).clamp_(0) - def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt): + def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt, image_size=None, epoch=0): """Select the positive anchor center in gt for rotated bounding boxes. Args: