# base_head

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

The base class every head inherits.

A head is a node with a task and an export parser. The task decides which losses, metrics, and visualizers can attach to the head.
The parser name goes into the NN Archive, together with the class names of the head.

## Classes

### BaseHead

Base class for all heads in the model.

A subclass sets the `task` class attribute. A subclass with an export parser also sets the `parser` class attribute. A subclass
can override
[get_custom_head_config](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_head.md)
to give its parser more values, and
[annotate](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_head.md)
to support a task that the default annotation does not know.

#### Methods

##### annotate

```python
def annotate(head_output: Packet[Tensor], image_paths: list[Path], config_preprocessing: PreprocessingConfig) -> DatasetIterator:
```

Convert the outputs of the head into dataset records.

This delegates to
[luxonis_train.utils.annotation.default_annotate](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/annotation.md).
Override it for tasks that the default converter does not support.

Parameters

 * `head_output` (`Packet[Tensor]`): The output packet of the head for one batch.
 * `image_paths` (`list[Path]`): The paths of the original images, in the order of the batch.
 * `config_preprocessing` (`PreprocessingConfig`): The preprocessing settings used to map predictions back to the images.

Returns

 * `DatasetIterator`: A generator of the annotation records.

##### get_custom_head_config

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

Return the head-specific metadata for the NN Archive.

A subclass overrides the method to give its parser more values.
[get_head_config](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_head.md)
merges the result into the `"metadata"` dictionary. The base implementation returns an empty dictionary.

Returns

 * `Params`: The additional metadata keys and their values.

##### get_head_config

```python
def get_head_config(self) -> dict[str, Any]:
```

Return the entry of the head in the NN Archive config.

The method starts from the `parser` class attribute, the
[BaseNode.class_names](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md),
and the
[BaseNode.n_classes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
of the head. Then it merges the result of
[get_custom_head_config](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_head.md)
into the `"metadata"` dictionary. A custom key replaces a base key of the same name.
[LuxonisModel.archive](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md)
calls the method for each head whose `remove_on_export` is `False`. Then it adds the `"name"` and the `"outputs"` keys.

The class names always come from `dataset_metadata`. Without it,
[BaseNode.class_names](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
raises `RuntimeError`. When the dataset has no task named `task_name`, it raises `ValueError`.

> **Example**
> ```pycon
>>> from torch import Size
>>> from luxonis_train.nodes import ClassificationHead
>>> from luxonis_train.utils import DatasetMetadata
>>> head = ClassificationHead(
...     task_name="animals",
...     dataset_metadata=DatasetMetadata(
...         classes={"animals": {"cat": 0, "dog": 1}}
...     ),
...     input_shapes=[{"features": [Size([1, 8, 4, 4])]}],
... )
>>> head.get_head_config()
{'parser': 'ClassificationParser',
 'metadata': {'classes': ['cat', 'dog'], 'n_classes': 2,
              'is_softmax': False}}
```

Returns

 * `dict[str, Any]`: A dictionary with the keys `"parser"` and `"metadata"`. The `"metadata"` dictionary holds `"classes"`, `"n_classes"`, and the custom keys.

#### Attributes

##### parser

The name of the parser that reads the outputs of the head in the exported model. [get_head_config](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/base_head.md) puts it into the NN Archive entry of the head. It is `""` in the base class.

##### task

The task of the head. It gives the key of the main output and the labels that the head needs.