# model

Python API: `luxonis_train.config.predefined_models.segmentation.v1.model`

The semantic segmentation model.

## Classes

### SegmentationModel

Semantic segmentation, which predicts a class for each pixel.

 * `Throughput:`: Frames per second at 384x512. * `light`: 43 on RVC2, 75 on RVC4
    * `heavy`: 14 on RVC2, 48 on RVC4

> **Notes**
> [JaccardIndex](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md) and [F1Score](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md) wrap `torchmetrics` and read a `task` key from `metrics_params`. Set it to `"binary"` for one class and to `"multiclass"` otherwise. Without `task`, a metric takes `"binary"` for one class and `"multiclass"` otherwise, and logs a warning.

By default, the model attaches an auxiliary head to an earlier output of the backbone. The loss of the auxiliary head adds to the
total loss with weight `0.4`. The export removes the auxiliary head.

> **Example**
> The `model` section of a config:

```yaml
model:
  predefined_model:
    name: SegmentationModel
    params:
      variant: light
```

 * `Components:`: * Nodes:
   [DDRNet](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/backbones/ddrnet/ddrnet.md)
   ->
   [DDRNetSegmentationHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/ddrnet_segmentation_head.md)
   ->
   [DDRNetSegmentationHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/ddrnet_segmentation_head.md)
    * Losses:
      [OHEMLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/ohem_loss.md)
    * Metrics: *
      [ConfusionMatrix](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/confusion_matrix/confusion_matrix.md)
       * [F1Score](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md)
       * [JaccardIndex](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md)
    * Visualizers:
      [SegmentationVisualizer](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md)
    * Main metric:
      [JaccardIndex](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md)
    * Variants: * `light`
       * `heavy`

#### Methods

##### init

```python
def __init__(use_aux_head: bool = True, aux_head_params: Params | None = None, **kwargs):
```

Initialize the model with its default components.

The constructor sets `attach_index` to `-1` in `head_params`, unless `head_params` holds an `aux_head` key. This value replaces an
`attach_index` that `head_params` sets.

Parameters

 * `use_aux_head` (`bool`): Add the auxiliary head to the node graph. See
   [nodes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/predefined_models/segmentation/v1/model.md).
 * `aux_head_params` (`Params | None`): The constructor parameters of the auxiliary head. The constructor edits a given non-empty
   dictionary in place: * It sets `attach_index` to `-2` when the dictionary does not hold the key. The auxiliary head then reads
   the second-to-last output of the backbone. For
   [DDRNet](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/backbones/ddrnet/ddrnet.md)
   with its default `use_aux_heads`, that output holds the high-resolution features.
    * It removes the `use_aux_heads` key, and uses its value as `remove_on_export` of the auxiliary head. Without the key, the
      value is `True`.
 * `**kwargs`: Keyword arguments for
   [SimplePredefinedModel.init](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/predefined_models/base_predefined_model.md).

Raises

 * `TypeError`: When the `use_aux_heads` value in `aux_head_params` is not a bool.

##### get_variants

```python
def get_variants() -> tuple[str, dict[str, Params]]:
```

Get the default variant name and the available variants.

The default is `light`. Both variants set `backbone` to `"DDRNet"` and `head` to `"DDRNetSegmentationHead"`. The `light` variant
uses the `"23-slim"` variant of
[DDRNet](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/backbones/ddrnet/ddrnet.md),
and `heavy` uses `"23"`.

Both variants set `weights` to `"download"` in `backbone_params` and in `head_params`, so the backbone and the head load their
COCO checkpoints. A `backbone_params` or `head_params` given in the config replaces the whole dictionary of the variant. The
auxiliary head loads no checkpoint.

> **Example**
> ```pycon
>>> default, variants = SegmentationModel.get_variants()
>>> default, variants["heavy"]["backbone_variant"]
('light', '23')
>>> variants["heavy"]["backbone_params"]
{'weights': 'download'}
```

Returns

 * `tuple[str, dict[str, Params]]`: `"light"` and the two variants with their constructor arguments.

#### Attributes

##### nodes

The backbone, the head, and the auxiliary head, as configs.

The first configs come from [SimplePredefinedModel.nodes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/predefined_models/base_predefined_model.md). When `use_aux_head` is `True`, the property appends a second config of the `head` class with these fields:

 * `alias`: the class name with the suffix `_aux`;
 * `inputs`: the backbone;
 * `params`: `aux_head_params`;
 * `task_name`: the `task_name` of the model;
 * `remove_on_export`: the `use_aux_heads` value of `aux_head_params`, `True` by default;
 * `losses`: the `loss` of the model with weight `0.4`, without `loss_params`.

The auxiliary head has no metrics, no visualizer, and no freezing. Its `variant` keeps the [NodeConfig](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/config.md) default `"default"`. [DDRNetSegmentationHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/ddrnet_segmentation_head.md) declares no variants, so the build of the node logs a warning and uses no variant parameters.

> **Example**
> ```pycon
>>> from luxonis_train.config.predefined_models import (
...     SegmentationModel,
... )
>>> model = SegmentationModel(variant="light")
>>> [(node.identifier, node.inputs) for node in model.nodes]
[('DDRNet', []),
 ('DDRNetSegmentationHead', ['DDRNet']),
 ('DDRNetSegmentationHead_aux', ['DDRNet'])]
>>> aux_head = model.nodes[-1]
>>> aux_head.params, aux_head.remove_on_export
({'attach_index': -2}, True)
>>> [(loss.name, loss.weight) for loss in aux_head.losses]
[('OHEMLoss', 0.4)]
```