# segmentation_head

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

The FCN segmentation head.

## Classes

### SegmentationHead

Basic FCN segmentation head with bilinear upsampling steps.

 * `Inputs:`: * `inputs` (`Tensor`): [B, C, H ⁄ s, W ⁄ s]
 * `Outputs:`: * `segmentation` (`Tensor`): [B, nclasses, H, W] logits. When s < 1, the logits keep the size of the feature map:
   [B, nclasses, H ⁄ s, W ⁄ s].

> **References**
> * Source: Adapted from [torchvision FCN](https://github.com/pytorch/vision/blob/main/torchvision/models/segmentation/fcn.py) (BSD-3-Clause).
 * License: [BSD-3-Clause](https://github.com/pytorch/vision/blob/main/LICENSE)

> **Notes**
> The stride s must be 2n, the same for the height and the width. The head applies n upsampling steps. Each step is an [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md): a bilinear upsample by `2`, a `1x1` convolution that halves the channels, and a `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) with batch norm and ReLU. A last `1x1` convolution gives one logit map for each class. The mode does not change the output key or shape. This includes export mode.

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

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

```yaml
- name: SegmentationHead
  inputs: [RepPANNeck]
```

 * `Compatible with:`: * Attach index: `-1`, the last output of the input node
    * Required labels: `segmentation`
    * 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`

#### Methods

##### init

```python
def __init__(**kwargs: Any):
```

Build the upsampling steps and the class convolution.

[infer_upscale_factor](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/general.md)
gives the number of steps n from the size of the feature map and the image size in `original_in_shape`. It raises `ValueError`
when a size ratio is not a power of two. It also raises it when the height ratio differs from the width ratio. Each step halves
the channels and rounds down. The last step thus has ⌊c ⁄ 2n⌋ channels, where c is the channel count of the feature map. When the
feature map is not smaller than the image, n is `0` or less. The head then has no upsampling step, and the logits keep the size of
the feature map.

The class annotates
[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_height](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md),
and
[BaseNode.in_width](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
as `int`. An `attach_index` of `"all"` or a range gives lists instead. Such an `attach_index` thus makes the constructor raise
[IncompatibleError](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/exceptions.md).

Parameters

 * `**kwargs` (`Any`): 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`, the input sizes through `input_shapes` or `in_sizes`, and the class count through
   `n_classes` or `dataset_metadata`.

##### forward

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

Compute the segmentation logits.

> **Example**
> A feature map with the stride `4` gets two upsampling steps:

```pycon
>>> import torch
>>> from torch import Size
>>> from luxonis_train.nodes import SegmentationHead
>>> head = SegmentationHead(
...     n_classes=3,
...     input_shapes=[{"features": [Size([1, 16, 8, 8])]}],
...     original_in_shape=Size([3, 32, 32]),
... )
>>> head(torch.zeros(1, 16, 8, 8)).shape
torch.Size([1, 3, 32, 32])
```

Parameters

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

Returns

 * `Tensor`: The logits of shape `[B, n_classes, H, W]`. When `s < 1`, the logits keep the size of the feature map, `[B,
   n_classes, H/s, W/s]`.
   [BaseNode.run](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/base_node.md)
   puts them 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.

#### Attributes

##### head

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