# atss_assigner

Python API: `luxonis_train.assigners.atss_assigner`

Adaptive Training Sample Selection, which selects the positive anchors of each ground truth box by an adaptive IoU threshold.

## Classes

### ATSSAssigner

Adaptive Training Sample Selection (ATSS) assigner.

The assigner selects the positive anchors of each ground truth box from the geometry of the anchors. It reads the predicted boxes
only to scale the assigned scores. For each ground truth box, it does these steps:

 * On each pyramid level, it selects up to `topk` anchors that have the centers closest to the center of the box. These anchors
   are the candidates.
 * It computes the IoU between each candidate and the box. The threshold of the box is μ + σ, the mean plus the standard deviation
   of these IoUs.
 * It keeps the candidates that have an IoU above the threshold and a center inside the box.

An anchor that stays positive for more than one box goes to the box that has the highest IoU with the anchor box. For details, see
[luxonis_train.assigners.utils.fix_collisions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/assigners/utils.md).

[AdaptiveDetectionLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/adaptive_detection_loss.md)
uses this assigner for the first `n_warmup_epochs` epochs. After these epochs, it uses
[TaskAlignedAssigner](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/assigners/tal_assigner.md).

> **References**
> * [Bridging the Gap Between Anchor-based and Anchor-free Detection via Adaptive Training Sample Selection](https://arxiv.org/pdf/1912.02424.pdf)
 * The implementation adapts the code of
   [PPYOLOE_pytorch](https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/assigner/atss_assigner.py) and
   [TOOD](https://github.com/fcjian/TOOD/blob/master/mmdet/core/bbox/assigners/atss_assigner.py).

#### Methods

##### init

```python
def __init__(n_classes: int, topk: int = 9):
```

Initialize the ATSS assigner.

Parameters

 * `n_classes` (`int`): Number of classes in the dataset. The label `n_classes` marks a background anchor in the output.
 * `topk` (`int`): Maximum number of candidate anchors to select on each pyramid level for each ground truth box. With fewer than
   three candidates for a box over all levels, no candidate passes the threshold.

##### forward

```python
def forward(anchor_bboxes: Tensor, n_level_bboxes: list[int], gt_labels: Tensor, gt_bboxes: Tensor, mask_gt: Tensor, pred_bboxes: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
```

Assign each anchor to a ground truth box or to the background.

All boxes are in `xyxy` format and in the same units. Only the boxes with `mask_gt` set to `1` get positive anchors.

> **Example**
> The box has an IoU of `0.64` with the first anchor box. The threshold is about `0.5`, so only the first anchor is positive. The predicted boxes equal the anchor boxes, so the score of that anchor is also `0.64`.

```pycon
>>> import torch
>>> assigner = ATSSAssigner(n_classes=2, topk=4)
>>> anchors = torch.tensor(
...     [
...         [0.0, 0.0, 4.0, 4.0],
...         [4.0, 0.0, 8.0, 4.0],
...         [0.0, 4.0, 4.0, 8.0],
...         [4.0, 4.0, 8.0, 8.0],
...     ]
... )
>>> gt_labels = torch.tensor([[[1.0]]])
>>> gt_bboxes = torch.tensor([[[0.0, 0.0, 5.0, 5.0]]])
>>> mask_gt = torch.tensor([[[1.0]]])
>>> labels, bboxes, scores, mask, gt_idx = assigner(
...     anchors, [4], gt_labels, gt_bboxes, mask_gt, anchors[None]
... )
>>> labels.tolist()
[[1, 2, 2, 2]]
>>> mask.tolist()
[[True, False, False, False]]
>>> [round(s, 2) for s in scores[0, 0].tolist()]
[0.0, 0.64]
```

Parameters

 * `anchor_bboxes` (`Tensor`): Anchor boxes with shape `[n_anchors, 4]`, ordered level by level.
 * `n_level_bboxes` (`list[int]`): Number of anchors on each pyramid level. The sum must equal `n_anchors`.
 * `gt_labels` (`Tensor`): Class index of each ground truth box with shape `[bs, n_max_boxes, 1]`.
 * `gt_bboxes` (`Tensor`): Ground truth boxes with shape `[bs, n_max_boxes, 4]`.
 * `mask_gt` (`Tensor`): `1` for a real box and `0` for a padded slot, with shape `[bs, n_max_boxes, 1]`.
 * `pred_bboxes` (`Tensor`): Predicted boxes with shape `[bs, n_anchors, 4]`. The IoU between a predicted box and its assigned box
   scales the assigned scores.

Returns

 * `tuple[Tensor, Tensor, Tensor, Tensor, Tensor]`: Five tensors. * `assigned_labels` (`[bs, n_anchors]`, `int64`) holds the class
   of the assigned box, or `n_classes` for a background anchor.
    * `assigned_bboxes` (`[bs, n_anchors, 4]`) holds the assigned box. Only the values at positive anchors are meaningful.
    * `assigned_scores` (`[bs, n_anchors, n_classes]`) holds a one-hot class vector scaled by the IoU between the predicted box
      and the assigned box. Zero for a background anchor.
    * `mask_positive` (`[bs, n_anchors]`, `bool`) is `True` at an anchor with an assigned box.
    * `assigned_gt_idx` (`[bs, n_anchors]`, `int64`) holds the index of the assigned box along dimension `1` of `gt_bboxes`, `0`
      for a background anchor. When `n_max_boxes` is `0`, every anchor is background. In this case, `mask_positive` and
      `assigned_gt_idx` are `float32` zeros.