# tal_assigner

Python API: `luxonis_train.assigners.tal_assigner`

The task-aligned assigner, which selects the positive anchors by the predicted class score and the IoU of the predicted box
together.

## Classes

### TaskAlignedAssigner

Task-aligned assigner (TAL) from TOOD.

The assigner scores each anchor for each ground truth box with the alignment metric:

t = sα⋅uβ

Here, s is the predicted score of the class of the box. The overlap u is the IoU between the predicted box of the anchor and the
ground truth box. With keypoints, u is this IoU times the object keypoint similarity from
[compute_pose_oks](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/keypoints.md).
For each ground truth box, the assigner does these steps:

 * It finds the anchors that have a center inside the box.
 * It keeps at most `topk` of these anchors, the ones with the highest t.

An anchor that stays positive for more than one box goes to the box with the highest u. 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).
The assigned score of a positive anchor is t, normalized for each box:

t̂ = (t)/(maxt)⋅maxu

Both maxima are over the positive anchors of the box.

Small-Target-Aware Label Assignment (STAL) helps small objects get positive anchors. When STAL is on, a side of a real box that is
shorter than the smallest stride gets the length of the second smallest stride. With only one stride, it gets the length of that
stride. The center of the box stays. The assigner uses the enlarged box only to find the anchors inside the box.

[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 after its warmup epochs.
[PrecisionDFLDetectionLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/precision_dfl_detection_loss.md)
uses it in all epochs.

> **References**
> * [TOOD: Task-aligned One-stage Object Detection](https://arxiv.org/pdf/2108.07755.pdf)
 * The implementation adapts the code of
   [PPYOLOE_pytorch](https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/assigner/tal_assigner.py), which has the
   [Apache License, Version 2.0](https://github.com/Nioolek/PPYOLOE_pytorch/tree/master?tab=Apache-2.0-1-ov-file#readme).

#### Methods

##### init

```python
def __init__(n_classes: int, topk: int = 13, alpha: float = 1.0, beta: float = 6.0, eps: float = 1e-09, strides: Sequence[int] | Tensor | None = None, skip_stal: bool = False):
```

Initialize the task-aligned assigner.

Parameters

 * `n_classes` (`int`): Number of classes in the dataset. The label `n_classes` marks a background anchor in the output.
 * `topk` (`int`): Number of anchors with the highest alignment metric that the assigner selects for each ground truth box. It
   must not be larger than the number of anchors.
 * `alpha` (`float`): The exponent α of the class score in the alignment metric.
 * `beta` (`float`): The exponent β of the overlap u in the alignment metric.
 * `eps` (`float`): A small value that prevents a division by zero in the score normalization and in the object keypoint
   similarity.
 * `strides` (`Sequence[int] | Tensor | None`): The strides of the detection head in pixels, for example `[8, 16, 32]`. The
   assigner sorts them and removes duplicates. STAL needs the strides, so `None` or an empty value turns STAL off.
 * `skip_stal` (`bool`): `True` turns STAL off. When `skip_stal` is `False` and `strides` is `None` or empty, the assigner logs a
   warning and turns STAL off.

##### forward

```python
def forward(pred_scores: Tensor, pred_bboxes: Tensor, anchor_points: Tensor, gt_labels: Tensor, gt_bboxes: Tensor, mask_gt: Tensor, pred_kpts: Tensor | None = None, gt_kpts: Tensor | None = None, sigmas: Tensor | None = None, area_factor: float | None = None) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
```

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

All boxes are in `xyxy` format. The boxes, the anchor points, and the keypoints use the same units. STAL compares the box sizes
with the strides, so STAL needs the boxes in pixels. Only the boxes with `mask_gt` set to `1` get positive anchors. The method
runs under `torch.no_grad`, so the outputs have no gradient.

To add the object keypoint similarity to the alignment metric, give all four of `pred_kpts`, `gt_kpts`, `sigmas`, and
`area_factor`.

> **Example**
> The box holds the centers of the first two anchors. Both predicted boxes have an IoU of `0.5` with it. The first one has the higher class score, so it gets the score `0.5`.

```pycon
>>> import torch
>>> assigner = TaskAlignedAssigner(
...     n_classes=2, topk=2, skip_stal=True
... )
>>> anchor_points = torch.tensor(
...     [[2.0, 2.0], [6.0, 2.0], [2.0, 6.0], [6.0, 6.0]]
... )
>>> pred_bboxes = torch.cat(
...     [anchor_points - 2, anchor_points + 2], -1
... )
>>> pred_scores = torch.tensor(
...     [[[0.1, 0.9], [0.2, 0.5], [0.3, 0.3], [0.9, 0.1]]]
... )
>>> gt_labels = torch.tensor([[[1.0]]])
>>> gt_bboxes = torch.tensor([[[0.0, 0.0, 8.0, 4.0]]])
>>> mask_gt = torch.tensor([[[1.0]]])
>>> labels, bboxes, scores, mask, gt_idx = assigner(
...     pred_scores,
...     pred_bboxes[None],
...     anchor_points,
...     gt_labels,
...     gt_bboxes,
...     mask_gt,
... )
>>> labels.tolist()
[[1, 1, 2, 2]]
>>> mask.tolist()
[[True, True, False, False]]
>>> [round(s, 2) for s in scores[0, :, 1].tolist()]
[0.5, 0.28, 0.0, 0.0]
```

Parameters

 * `pred_scores` (`Tensor`): Predicted class probabilities in `[0, 1]` with shape `[bs, n_anchors, n_classes]`. The losses give
   the sigmoid of the class logits.
 * `pred_bboxes` (`Tensor`): Predicted boxes with shape `[bs, n_anchors, 4]`.
 * `anchor_points` (`Tensor`): Anchor centers `(x, y)` with shape `[n_anchors, 2]`.
 * `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_kpts` (`Tensor | None`): Predicted keypoints with shape `[bs, n_anchors, n_kpts, 3]`. The assigner reads only `x` and
   `y`.
 * `gt_kpts` (`Tensor | None`): Ground truth keypoints as `(x, y, visibility)` with shape `[bs, n_max_boxes, n_kpts, 3]`.
 * `sigmas` (`Tensor | None`): One sigma per keypoint with shape `[n_kpts]`.
 * `area_factor` (`float | None`): The factor that scales the area of a ground truth box to the pose area.

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 normalized alignment metric t̂.
      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.

Raises

 * `ValueError`: When some, but not all, of `pred_kpts`, `gt_kpts`, `sigmas`, and `area_factor` are `None`.