# ddrnet_segmentation_head

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

The segmentation head of DDRNet.

## Classes

### DDRNetSegmentationHead

DDRNet segmentation head.

 * `Inputs:`: * `inputs` (`Tensor`): [B, C, H ⁄ s, W ⁄ s]
 * `Outputs:`: * train, eval: * `segmentation` (`Tensor`): [B, nclasses, H, W] logits
    * export: * `segmentation` (`Tensor`): [B, H, W] `int32` class indices, or [B, 1, H, W] with the values `0` and `1` when
      `n_classes == 1`

> **References**
> * Source: Adapted from [Deci-AI/super-gradients](https://github.com/Deci-AI/super-gradients) (Apache-2.0).
 * License: Apache-2.0

> **Notes**
> H and W are the height and the width of the model input. The scale s is a power of two that the head computes from the input size and the model input size. The [DDRNet](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/backbones/ddrnet/ddrnet.md) backbone gives s = 8.

The head applies batch norm and ReLU, a 3x3 convolution, batch norm and ReLU, and a 1x1 convolution to nclasses channels. An
upsampling layer then scales the result by s. In export mode, the head converts the logits to class indices.

Warning: With `inter_mode="pixel_shuffle"`, the 1x1 convolution keeps `inter_channels` channels. The number of output channels is
then `inter_channels` divided by s2, not nclasses.

 * `Variants:`: None. Configure the node through `params`.

> **See Also**
> * [DDRNet in super-gradients](https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/models/segmentation_models/ddrnet.py)
 * [The original DDRNet code](https://github.com/ydhongHIT/DDRNet)
 * [The DDRNet paper](https://arxiv.org/pdf/2101.06085.pdf)

> **Example**
> A node entry in the `model.nodes` section of a config:

```yaml
- name: DDRNetSegmentationHead
  inputs: [DDRNet]
```

 * `Compatible with:`: * Attach index: `-1`, the last output of the input node
    * Required labels: `segmentation`
    * Used by:
      [SegmentationModel](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/predefined_models/segmentation/v1/model.md)
    * Losses: *
      [BCEWithLogitsLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/bce_with_logits.md)
       * [CrossEntropyLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/cross_entropy.md)
       * [OHEMLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/ohem_loss.md)
       * [SigmoidFocalLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/sigmoid_focal_loss.md)
       * [SmoothBCEWithLogitsLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/smooth_bce_with_logits.md)
       * [SoftmaxFocalLoss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/losses/softmax_focal_loss.md)
    * Metrics: *
      [Accuracy](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md)
       * [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)
       * [DiceCoefficient](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/dice_coefficient.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)
       * [MIoU](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/mean_iou.md)
       * [Precision](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/metrics/torchmetrics.md)
       * [Recall](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)
    * Export parser: `SegmentationParser`
    * Pretrained weights: available through `weights: download`

#### Methods

##### init

```python
def __init__(inter_channels: int = 64, inter_mode: Literal['nearest', 'linear', 'bilinear', 'bicubic', 'trilinear', 'area', 'pixel_shuffle'] = 'bilinear', **kwargs):
```

Build the layers and the upsampling of the head.

The constructor computes the scale s = 2n.
[infer_upscale_factor](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/general.md)
gives n from the input size and the model input size. That function raises `ValueError` when the height ratio or the width ratio
is not a power of two, or when the two ratios differ.

Parameters

 * `inter_channels` (`int`): The number of output channels of the 3x3 convolution. With `"pixel_shuffle"`, it must be a multiple
   of s2.
 * `inter_mode` (`Literal['nearest', 'linear', 'bilinear', 'bicubic', 'trilinear', 'area', 'pixel_shuffle']`): The upsampling
   method. `"pixel_shuffle"` uses `torch.nn.PixelShuffle`. The other values are the `mode` of `torch.nn.Upsample`. With `"linear"`
   or `"trilinear"`,
   [forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/ddrnet_segmentation_head.md)
   raises `NotImplementedError`, because these modes need a 3D or a 5D input.
 * `**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).
   They must hold `original_in_shape`, `input_shapes` or `in_sizes`, and the class count through `n_classes` or
   `dataset_metadata`.

Raises

 * `ValueError`: When `inter_mode` is `"pixel_shuffle"` and `inter_channels` is not a multiple of s2.

##### forward

```python
def forward(inputs: Tensor) -> Tensor:
```

Compute the logits, or the class indices in export mode.

> **Example**
> ```pycon
>>> import torch
>>> from torch import Size
>>> from luxonis_train.nodes import DDRNetSegmentationHead
>>> head = DDRNetSegmentationHead(
...     n_classes=3,
...     input_shapes=[{"features": [Size([1, 8, 4, 4])]}],
...     original_in_shape=Size([3, 32, 32]),
... )
>>> head.eval()(torch.zeros(1, 8, 4, 4)).shape
torch.Size([1, 3, 32, 32])
>>> head.export = True
>>> out = head(torch.zeros(1, 8, 4, 4))
>>> out.shape, out.dtype
(torch.Size([1, 32, 32]), torch.int32)
```

Parameters

 * `inputs` (`Tensor`): The feature map of shape `[B, C, H / s, W / s]`.

Returns

 * `Tensor`: Outside export mode, the logits of shape `[B, n_classes, H, W]`. In export mode, an `int32` tensor. It holds the class index with the highest logit for each pixel, of shape `[B, H, W]`. With `n_classes == 1`, it holds `1` where the logit is above `0`, of shape `[B, 1, H, W]`. [BaseNode.run](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md) puts the result under the `"segmentation"` key.

##### 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_weights_url

```python
def get_weights_url(self) -> str:
```

Select the COCO checkpoint from the input channel count.

The 128-channel `23-slim` backbone uses the slim checkpoint; the 256-channel `23` backbone uses the full checkpoint.

Raises

 * `NotImplementedError`: If the input has neither 128 nor 256 channels.

##### load_checkpoint

```python
def load_checkpoint(path: str | None = None, strict: bool = False):
```

Load a checkpoint, with a non-strict key match by default.

The method passes `path` as `ckpt` to [BaseNode.load_checkpoint](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md). The override renames the first parameter and sets the default of `strict` to `False`.

Warning: The override has no `ckpt` parameter. After construction, the node calls `load_checkpoint(ckpt=...)` for a `weights` URL. Thus, a call such as `DDRNetSegmentationHead(weights="https://...")` raises `TypeError`. For `weights="download"`, the node calls the method without arguments, so that value does not fail this way.

Parameters

 * `path` (`str | None`): Local path or URL of a `.ckpt` file. [LuxonisLightningModule](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/lightning/luxonis_lightning.md) also gives a state dictionary, and the base method loads it directly. `None` or `""` takes the URL from [get_weights_url](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/ddrnet_segmentation_head.md).
 * `strict` (`bool`): Whether the keys of the checkpoint must match the keys of the head exactly.

#### Attributes

##### bn1

##### bn2

##### conv1

##### conv2

##### 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_height

The height of the attached inputs.

It is the second 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). A list of sizes gives a list of heights.

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_width

The width of the attached inputs.

It is the last dimension of [in_sizes](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md). A list of sizes gives a list of widths.

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.

##### parser

##### relu

##### upscale