# segmentation_visualizer

Python API: `luxonis_train.attached_modules.visualizers.segmentation_visualizer`

The visualizer that blends segmentation masks into the images.

## Classes

### SegmentationVisualizer

Visualizer for semantic segmentation and anomaly masks.

The left image shows the targets. The right image shows the predictions.

 * `Inputs:`: * `prediction_canvas`, `target_canvas` (`Tensor`): [B, 3, H, W]
    * `predictions` (`Tensor`): [B, nclasses, H, W] logits
    * `target` (`Tensor | None`): [B, nclasses, H, W] one-hot
 * `Outputs:`: * `Tensor | tuple[Tensor, Tensor]`: [B, 3, H, W], a pair when targets are given

> **References**
> * Source: This project.
 * License: Apache-2.0 (this project)

> **Notes**
> Each predicted pixel gets the class with the highest logit. With one class, a pixel gets the class when the sigmoid of its logit is at least `0.5`. The visualizer blends one color for each class into the images. On an `anomaly_detection` node, it reads the `segmentation` label.

> **Example**
> Attached to a `DDRNetSegmentationHead` in `model.nodes`:

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

 * `Compatible with:`: * Used by: *
   [AnomalyDetectionModel](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/predefined_models/anomaly_detection/v1/model.md)
       * [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)
    * Nodes: *
      [BiSeNetHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/bisenet_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)
       * [DiscSubNetHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/discsubnet_head/discsubnet_head.md)
       * [SegmentationHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/segmentation_head.md)
       * [TransformerSegmentationHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/transformer_segmentation_head.md)

#### Methods

##### init

```python
def __init__(colors: Color | list[Color] | None = None, background_class: int | None = 0, background_color: Color = '#000000', alpha: float = 0.6, **kwargs):
```

Initialize the visualizer and store the color options.

Parameters

 * `colors` (`Color | list[Color] | None`): One color for each class, in the order of the class indices. A single color becomes a
   list of one color. When `None`, or when the number of colors is not the number of classes,
   [forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md)
   uses the colors of
   [BaseVisualizer.colormap](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/base_visualizer.md).
   It logs a warning on its first call.
 * `background_class` (`int | None`): The index of the class that gets `background_color`. It applies only when
   [forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md)
   uses the colormap colors and the node has more than one class. `None` gives every class a colormap color.
 * `background_color` (`Color`): The color of the background class.
 * `alpha` (`float`): The opacity of the masks, from `0` for transparent to `1` for opaque.
 * `**kwargs`: Keyword arguments forwarded to
   [BaseVisualizer](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/base_visualizer.md),
   such as `scale` and `node`.

##### draw_predictions

```python
def draw_predictions(canvas: Tensor, predictions: Tensor, alpha: float, colors: list[Color], scale: float = 1.0) -> Tensor:
```

Draw the predicted masks of a batch on copies of the canvas.

For each image,
[seg_output_to_bool](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/utils/segmentation.md)
converts the logits to one boolean mask for each class. With more than one class, each pixel gets the class with the highest
logit. With one class, a pixel gets the class when the sigmoid of its logit is at least `0.5`.
[potentially_upscale_masks](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/utils.md)
resizes the masks by `scale`, and
[draw_segmentation_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/utils.md)
blends them into the image.

> **Example**
> The first pixel column is class `0` and the second is class `1`. The masks double in size.

```pycon
>>> import torch
>>> canvas = torch.zeros(1, 3, 2, 4, dtype=torch.uint8)
>>> logits = torch.tensor([[[[2.0, 0.0]], [[0.0, 2.0]]]])
>>> viz = SegmentationVisualizer.draw_predictions(
...     canvas,
...     logits,
...     alpha=1.0,
...     colors=["red", "blue"],
...     scale=2.0,
... )
>>> viz[0, 0].tolist()
[[255, 255, 0, 0], [255, 255, 0, 0]]
```

Parameters

 * `canvas` (`Tensor`): `uint8` images of shape `[B, 3, H, W]`. The method does not change them.
 * `predictions` (`Tensor`): Logits of shape `[B, n_classes, h, w]`. The masks must have the canvas size after the resize by
   `scale`.
 * `alpha` (`float`): The opacity of the masks, from `0` to `1`.
 * `colors` (`list[Color]`): One color for each class, at least as many colors as classes.
 * `scale` (`float`): The factor that resizes the masks.

Returns

 * `Tensor`: A new tensor of the same shape as `canvas`, with the masks drawn.

##### draw_targets

```python
def draw_targets(canvas: Tensor, targets: Tensor, alpha: float, colors: list[Color], scale: float = 1.0) -> Tensor:
```

Draw the target masks of a batch on copies of the canvas.

For each image, the method casts the target to `bool`, so every non-zero value marks a pixel of the class.
[potentially_upscale_masks](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/utils.md)
resizes the masks by `scale`, and
[draw_segmentation_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/utils.md)
blends them into the image.

> **Example**
> ```pycon
>>> import torch
>>> canvas = torch.zeros(1, 3, 2, 2, dtype=torch.uint8)
>>> target = torch.tensor([[[[1, 0], [0, 0]], [[0, 1], [1, 1]]]])
>>> viz = SegmentationVisualizer.draw_targets(
...     canvas, target, alpha=1.0, colors=["red", "blue"]
... )
>>> viz[0, :, 0, 0].tolist(), viz[0, :, 0, 1].tolist()
([255, 0, 0], [0, 0, 255])
```

Parameters

 * `canvas` (`Tensor`): `uint8` images of shape `[B, 3, H, W]`. The method does not change them.
 * `targets` (`Tensor`): One-hot masks of shape `[B, n_classes, h, w]`. The masks must have the canvas size after the resize by `scale`.
 * `alpha` (`float`): The opacity of the masks, from `0` to `1`.
 * `colors` (`list[Color]`): One color for each class, at least as many colors as classes.
 * `scale` (`float`): The factor that resizes the masks.

Returns

 * `Tensor`: A new tensor of the same shape as `canvas`, with the masks drawn.

##### forward

```python
def forward(prediction_canvas: Tensor, target_canvas: Tensor, predictions: Tensor, target: Tensor | None) -> tuple[Tensor, Tensor]
| Tensor:
```

Draw the predicted masks, and the target masks when given.

The method first selects the class colors. It uses `colors` when the list holds one color for each class of the node. Otherwise it takes the colors of [BaseVisualizer.colormap](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/base_visualizer.md), and logs a warning on the first call. With more than one class, the `background_class` then gets the `background_color`. [draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md) and [draw_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md) resize the masks by the `scale` factor and draw them.

Parameters

 * `prediction_canvas` (`Tensor`): `uint8` images of shape `[B, 3, H, W]` to draw the predictions on.
 * `target_canvas` (`Tensor`): `uint8` images of shape `[B, 3, H, W]` to draw the targets on.
 * `predictions` (`Tensor`): Logits of shape `[B, n_classes, h, w]`, the main output of the node.
 * `target` (`Tensor | None`): One-hot masks of shape `[B, n_classes, h, w]`, the `segmentation` label. `None` when the batch has no such label.

Returns

 * `tuple[Tensor, Tensor] | Tensor`: The predictions image when `target` is `None`, otherwise the pair `(targets, predictions)`. Each image has the shape of its canvas.

#### Attributes

##### required_labels

The labels for the `target` parameter of [forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md).

[BaseAttachedModule.get_parameters](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/base_attached_module.md) reads this set for the `target` parameter of [forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/segmentation_visualizer.md), which has no label suffix. On an `anomaly_detection` node, the set holds only `"segmentation"`, so `target` receives the anomaly mask. On other nodes, it holds the labels of the task. The property raises `RuntimeError` when the visualizer has no task.

##### supported_tasks

## Attributes

### log_disable