# base_detection_head

Python API: `luxonis_train.nodes.heads.base_detection_head`

The base class of the heads that predict bounding boxes.

It keeps the NMS settings and the strides of the scales. It also selects the names of the exported outputs. Each subclass runs NMS
itself, in evaluation mode.

## Classes

### BaseDetectionHead

Base class for YOLO-like detection heads with several scales.

The head reads the last `n_heads` outputs of the input node, one feature map for each scale. For this selection, the constructor
sets `attach_index` to `(-n_heads - 1, -1)`. An `attach_index` in the node `params` replaces this value.

A subclass stores one block for each scale in `heads`. In evaluation mode, it runs NMS with `conf_thres`, `iou_thres`, and
`max_det`. After a call to
[request_detections_pre_nms](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_detection_head.md),
it also keeps the NMS input in its packet.

> **Example**
> [EfficientBBoxHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/efficient_bbox_head.md) is a detection head. With two scales, it reads the last two of three feature maps:

```pycon
>>> from torch import Size
>>> from luxonis_train.nodes import EfficientBBoxHead
>>> sizes = [
...     Size([1, 4, 64, 64]),
...     Size([1, 8, 32, 32]),
...     Size([1, 16, 16, 16]),
... ]
>>> head = EfficientBBoxHead(
...     n_heads=2,
...     n_classes=3,
...     input_shapes=[{"features": sizes}],
...     original_in_shape=Size([3, 256, 256]),
... )
>>> head.attach_index, head.in_channels, head.stride.tolist()
((-3, -1), [8, 16], [8, 16])
```

#### Methods

##### init

```python
def __init__(n_heads: int, conf_thres: float, iou_thres: float, max_det: int, **kwargs):
```

Set the NMS settings, the attach index, and the strides.

The constructor reads
[BaseNode.in_channels](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md),
[BaseNode.in_sizes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md),
and
[BaseNode.original_in_shape](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md).
Thus `kwargs` must hold `original_in_shape`, and `input_shapes` or `in_sizes`. When the head gets fewer than `n_heads` feature
maps, the constructor logs a warning and sets `n_heads` to that number. Without an `attach_index` in `kwargs`, this check counts
all outputs of the input node.

Parameters

 * `n_heads` (`int`): The number of scales. The head reads the last `n_heads` outputs of the input node.
 * `conf_thres` (`float`): The confidence threshold of NMS, in `[0, 1]`.
 * `iou_thres` (`float`): The IoU threshold of NMS, in `[0, 1]`.
 * `max_det` (`int`): The maximum number of boxes that NMS keeps for each image.
 * `**kwargs`: Keyword arguments for
   [BaseNode](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md).
   An `attach_index` among them replaces the default selection of the last `n_heads` outputs. It must select a range or `"all"`.
   An integer index makes the constructor fail. This class annotates `in_channels` as `list[int]`, so
   [BaseNode](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
   raises
   [IncompatibleError](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/exceptions.md).
   A subclass with its own class annotations hides this annotation, for example
   [PrecisionSegmentBBoxHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/precision_seg_bbox_head.md).
   For such a subclass, the constructor raises `TypeError` instead.

##### fit_stride_to_heads

```python
def fit_stride_to_heads(self) -> Tensor:
```

Compute the stride of each scale from the input sizes.

The stride of scale i is si = round(H ⁄ Hi), where H is the height of the model input and Hi is the height of the feature map. The
method reads the first `n_heads` sizes of
[BaseNode.in_sizes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md).
It takes Hi from index `2` of each size, so the sizes must have the form `[B, C, H, W]`. The constructor stores the result in
`stride`. The class example shows the result.

Returns

 * `Tensor`: An `int32` tensor of shape `[n_heads]`.

##### get_custom_head_config

```python
def get_custom_head_config(self) -> Params:
```

Return the NMS settings and the strides for the NN Archive.

A subclass adds its own keys to this dictionary, for example `"subtype"`.

Returns

 * `Params`: A dictionary with the keys `"iou_threshold"`, `"conf_threshold"`, `"max_det"`, and `"strides"`. They hold
   `iou_thres`, `conf_thres`, `max_det`, and `stride` as a list with one integer for each scale.

##### get_output_names

```python
def get_output_names(default: list[str]) -> list[str]:
```

Return the export output names or the `default` names.

A subclass calls the method in its `export_output_names` property. The method reads the `export_output_names` constructor argument
through
[BaseNode.export_output_names](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md).
It does not read the property of the subclass, which calls this method. It returns the names of the constructor argument when
their number is `n_heads`. Otherwise, it logs a warning and returns `default`. It also logs a warning when the argument is `None`.

Warning: The method compares the number of names with `n_heads`, not with the length of `default`. A subclass with more than
`n_heads` outputs thus never gets names for all its outputs from the argument.

> **Example**
> ```pycon
>>> from torch import Size
>>> from luxonis_train.nodes import EfficientBBoxHead
>>> sizes = [Size([1, 8, 32, 32]), Size([1, 16, 16, 16])]
>>> def build_head(names):
...     return EfficientBBoxHead(
...         n_heads=2,
...         n_classes=3,
...         input_shapes=[{"features": sizes}],
...         original_in_shape=Size([3, 256, 256]),
...         export_output_names=names,
...     )
>>> build_head(["small", "large"]).get_output_names(["a", "b"])
['small', 'large']
```

One name for two scales gives the default names. The example turns the logger off, so the warning does not show:

```pycon
>>> from loguru import logger
>>> logger.disable("luxonis_train")
>>> build_head(["boxes"]).get_output_names(["a", "b"])
['a', 'b']
>>> logger.enable("luxonis_train")
```

Parameters

 * `default` (`list[str]`): The names to return when the constructor argument is `None` or has the wrong length. The subclasses give names that DepthAI accepts.

Returns

 * `list[str]`: The names of the constructor argument, or `default`.

##### request_detections_pre_nms

```python
def request_detections_pre_nms(self):
```

Make the head add the pre-NMS candidates to its packet.

After the call, the evaluation packet of a subclass also holds the `"detections_pre_nms"` key. Its value is the NMS input: a tensor of shape `[B, N, 5 + n_classes]` with one row for each of the `N` anchor points. Each row holds the `xyxy` box in pixels, a constant `1`, and the class scores. A head that keeps more values with each box, such as keypoints or mask coefficients, adds columns after the class scores. The tensor can use much memory, so the head keeps it only on request. [PrecisionRecallCurve](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/precision_recall_curve.md) calls this method in its constructor.

> **Example**
> ```pycon
>>> import torch
>>> from torch import Size
>>> from luxonis_train.nodes import EfficientBBoxHead
>>> sizes = [Size([1, 8, 32, 32]), Size([1, 16, 16, 16])]
>>> head = EfficientBBoxHead(
...     n_heads=2,
...     n_classes=3,
...     input_shapes=[{"features": sizes}],
...     original_in_shape=Size([3, 256, 256]),
... )
>>> head.keep_detections_pre_nms
False
>>> head.request_detections_pre_nms()
>>> out = head.eval()([torch.zeros(size) for size in sizes])
>>> out["detections_pre_nms"].shape
torch.Size([1, 1280, 8])
```

#### Attributes

##### attach_index

##### in_channels

The number of channels of the attached inputs.

It is the third dimension from the end of
[in_sizes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md),
so a shape with or without the batch dimension gives the same value. A list of sizes gives a list of channel counts.

Raises

 * `RuntimeError`: When
   [in_sizes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
   cannot find the input sizes.
 * `ValueError`: When
   [attach_index](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
   does not fit the sizes.

##### in_sizes

The sizes of the attached inputs.

The property uses the first rule that applies:

 1. The `in_sizes` constructor argument, when it is set.
 2. The `"features"` entry of the only input packet.
 3. The only entry of that packet.
 4. The entries of that packet whose keys match the names of the `forward` parameters. Their sizes must be equal, and the property
    uses the first one.

Rules 2 to 4 pass the entry through
[get_attached](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md).
The result is a single size for an integer
[attach_index](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md),
and a list of sizes for `"all"` or a range. A node with more than one input, or with shapes that the rules do not fit, must read
[input_shapes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
instead.

> **Example**
> ```pycon
>>> from torch import Size, Tensor
>>> from luxonis_train.nodes import BaseNode
>>> class Node(BaseNode, register=False):
...     def forward(self, x: list[Tensor]) -> list[Tensor]:
...         return x
>>> shapes = [
...     {"features": [Size([2, 8, 64, 64]), Size([2, 16, 32, 32])]}
... ]
>>> node = Node(input_shapes=shapes)
>>> node.attach_index
'all'
>>> node.in_sizes
[torch.Size([2, 8, 64, 64]), torch.Size([2, 16, 32, 32])]
>>> node.in_channels, node.in_height, node.in_width
([8, 16], [64, 32], [64, 32])
```

Raises

 * `RuntimeError`: When `input_shapes` is missing or does not hold exactly one packet. Also when no key matches a `forward` parameter, or when the matching sizes differ. Also when [attach_index](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md) is `None` and the entry is a list.
 * `ValueError`: When [attach_index](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md) does not fit the sizes.

##### iou_thres

The IoU threshold of NMS.

##### keep_detections_pre_nms

Whether the head adds the pre-NMS candidates to its packet.

It is `False` until a call to [request_detections_pre_nms](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_detection_head.md). When it is `True`, the evaluation packet of a subclass also holds the `"detections_pre_nms"` key.

##### max_det

The maximum number of boxes that NMS keeps for each image.

##### parser

The export parser, `"YOLO"`. A subclass can replace it.

##### stride

The stride of each scale, an `int32` tensor of shape `[n_heads]`. [fit_stride_to_heads](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_detection_head.md) computes it.