# keypoint_visualizer

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

Draws keypoints, the skeleton that connects them, and the boxes they belong to.

## Classes

### KeypointVisualizer

Visualizer for instance keypoints and their bounding boxes.

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

 * `Inputs:`: * `prediction_canvas`, `target_canvas` (`Tensor`): [B, 3, H, W]
    * `keypoints` (`list[Tensor]`): [Mi, nkeypoints, 3] per image, `(x, y, conf)`, pixels
    * `boundingbox` (`list[Tensor]`): [Mi, 6] per image, `[x1, y1, x2, y2, conf, class]`, pixels
    * `target_keypoints` (`Tensor | None`): [N, 1 + 3**n*keypoints], `[batch, x, y, v, ...]`, normalized
    * `target_boundingbox` (`Tensor | None`): [N, 6], `[batch, class, x, y, w, h]`, `xywh` normalized
 * `Outputs:`: * `Tensor | tuple[Tensor, Tensor]`: [B, 3, H, W], a `(targets, predictions)` pair when any target is given

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

> **Notes**
> Draws the boxes with [BBoxVisualizer](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/bbox_visualizer.md), then the keypoints, the optional skeleton lines, and the optional keypoint indices on top. The `FOMO` task is in `supported_tasks`, but [FOMOHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/fomo_head.md) puts no `boundingbox` key in its packet. On that node, [BaseVisualizer.run](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/base_visualizer.md) raises `RuntimeError`.

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

```yaml
- name: EfficientKeypointBBoxHead
  inputs: [RepPANNeck]
  visualizers:
    - name: KeypointVisualizer
```

 * `Compatible with:`: * Used by:
   [KeypointDetectionModel](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/config/predefined_models/keypoint_detection/v1/model.md)
    * Nodes: *
      [EfficientKeypointBBoxHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/efficient_keypoint_bbox_head.md)
       * [FOMOHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/fomo_head.md)

#### Methods

##### init

```python
def __init__(visibility_threshold: float = 0.5, connectivity: list[tuple[int, int]] | None = None, visible_color: Color = 'red', nonvisible_color: Color | None = None, radius: int | None = None, draw_indices: bool = False, **kwargs):
```

Initialize the visualizer and store the keypoint options.

Parameters

 * `visibility_threshold` (`float`): The lowest confidence of a visible predicted keypoint.
   [draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md)
   tells how the visualizer draws the other keypoints.
 * `connectivity` (`list[tuple[int, int]] | None`): Pairs of keypoint indices to connect with lines, the skeleton. Applies to the
   predictions and the targets. `None` draws no lines.
 * `visible_color` (`Color`): Color of the visible predicted keypoints, and of all target keypoints. A color name such as `"red"`
   or an RGB tuple.
 * `nonvisible_color` (`Color | None`): Color of the predicted keypoints below `visibility_threshold`. When `None`, the visualizer
   does not draw them at their coordinates.
 * `radius` (`int | None`): Radius of a keypoint, in pixels. When `None`,
   [forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md)
   picks it from the size of each canvas.
 * `draw_indices` (`bool`): Whether to write the index of each keypoint next to it.
   [draw_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md)
   tells when this raises `RuntimeError` for the target keypoints.
 * `**kwargs`: Keyword arguments forwarded to
   [BBoxVisualizer](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/bbox_visualizer.md),
   such as `labels`, `colors`, `width`, `scale`, and `node`.

##### draw_keypoint_indices_pil

```python
def draw_keypoint_indices_pil(canvas: Tensor, keypoints: Tensor, offset: tuple[int, int] = (7, 7), colors: Color = 'red') -> Tensor:
```

Write the index of each keypoint next to it with PIL.

The method flattens `keypoints` to rows of `(x, y, visibility)` and numbers the rows from `0` in one sequence. With several
instances, the numbers do not restart for each instance. It does not read the visibility.

The method centers each label on its keypoint and then shifts it by `offset`. The direction of the shift cycles from one row to
the next: down-left, down-right, up-right, and up-left.

> **Example**
> ```pycon
>>> import torch
>>> canvas = torch.zeros(3, 32, 32, dtype=torch.uint8)
>>> keypoints = torch.tensor([[16.0, 16.0, 1.0]])
>>> viz = KeypointVisualizer.draw_keypoint_indices_pil(
...     canvas, keypoints
... )
>>> viz.dtype, viz.shape
(torch.float32, torch.Size([3, 32, 32]))
>>> bool(viz[0].any()), bool(viz[1].any())
(True, False)
```

Parameters

 * `canvas` (`Tensor`): One `uint8` image of shape `[3, H, W]`. PIL raises `TypeError` for a floating point image.
 * `keypoints` (`Tensor`): Pixel keypoints with three values per keypoint, such as `[M, K, 3]` or `[M, 3 * K]`. The method calls `view`, so a tensor that it cannot view as `[-1, 3]` raises `RuntimeError`.
 * `offset` (`tuple[int, int]`): The vertical and the horizontal shift of a label, in pixels.
 * `colors` (`Color`): Text color.

Returns

 * `Tensor`: A new `float32` image of shape `[3, H, W]` on the CPU, with values from `0` to `255` and the indices drawn.

##### draw_predictions

```python
def draw_predictions(canvas: Tensor, predictions: list[Tensor], draw_indices: bool = False, nonvisible_color: Color | None = None,
visible_color: Color = 'red', visibility_threshold: float = 0.5, radius: int | None = None, scale: float = 1.0, **kwargs) ->
Tensor:
```

Draw the predicted keypoints of a batch on copies of the canvas images.

For each image, the method multiplies the coordinates by `scale`. A keypoint with a confidence below `visibility_threshold` is not visible. Then the method draws the keypoints with `torchvision.utils.draw_keypoints` in two passes:

 * The first pass draws the visible keypoints in `visible_color`, clamped into the image. It moves each keypoint that is not visible to the top-left corner `(0, 0)` and draws it there too.
 * The second pass runs only when `nonvisible_color` is set. It draws the keypoints that are not visible in that color, and does not clamp them. It moves each visible keypoint to `(0, 0)` and draws it there too.

When `draw_indices` is set, each pass also writes the keypoint indices at the positions it uses, with [draw_keypoint_indices_pil](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md), in the color of the pass.

> **Examples**
> A visible keypoint stays at its coordinates:

```pycon
>>> import torch
>>> canvas = torch.zeros(1, 3, 16, 16, dtype=torch.uint8)
>>> keypoints = [torch.tensor([[[8.0, 8.0, 0.9]]])]
>>> viz = KeypointVisualizer.draw_predictions(
...     canvas, keypoints, radius=1
... )
>>> viz[0, :, 8, 8].tolist(), viz[0, :, 0, 0].tolist()
([255, 0, 0], [0, 0, 0])
```

A keypoint below `visibility_threshold` moves to the top-left corner:

```pycon
>>> hidden = [torch.tensor([[[8.0, 8.0, 0.1]]])]
>>> viz = KeypointVisualizer.draw_predictions(
...     canvas, hidden, radius=1
... )
>>> viz[0, :, 8, 8].tolist(), viz[0, :, 0, 0].tolist()
([0, 0, 0], [255, 0, 0])
```

Parameters

 * `canvas` (`Tensor`): `uint8` images of shape `[B, 3, H, W]`. The method does not modify it.
 * `predictions` (`list[Tensor]`): One tensor per image, of shape `[M_i, K, 3]`. Each keypoint is `(x, y, confidence)`. The coordinates are pixels of the unscaled image.
 * `draw_indices` (`bool`): Whether to write the index of each keypoint next to it.
 * `nonvisible_color` (`Color | None`): Color of the second pass. `None` skips the second pass.
 * `visible_color` (`Color`): Color of the first pass. A `colors` key in `kwargs` replaces it for the keypoints, but not for the indices.
 * `visibility_threshold` (`float`): The lowest confidence of a visible keypoint.
 * `radius` (`int | None`): Radius of a keypoint, in pixels. When it is `None`, `torchvision` raises `TypeError` for an image with at least one keypoint.
 * `scale` (`float`): Multiplier for the coordinates. Pass the factor that scaled the canvas.
 * `**kwargs`: Keyword arguments for `torchvision.utils.draw_keypoints`, such as `connectivity` and `width`.

Returns

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

##### draw_targets

```python
def draw_targets(canvas: Tensor, targets: Tensor, draw_indices: bool = False, colors: Color = 'red', **kwargs) -> Tensor:
```

Draw the target keypoints of a batch on copies of the canvas images.

The keypoints of image `i` are the rows of `targets` whose first column equals `i`. [draw_keypoint_labels](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/utils.md) converts them from normalized to pixel coordinates with the canvas size and draws them with `torchvision.utils.draw_keypoints`. It draws every keypoint, whatever its visibility.

> **Example**
> ```pycon
>>> import torch
>>> canvas = torch.zeros(1, 3, 16, 16, dtype=torch.uint8)
>>> targets = torch.tensor([[0, 0.5, 0.5, 2.0]])
>>> viz = KeypointVisualizer.draw_targets(
...     canvas, targets, radius=1
... )
>>> viz[0, :, 8, 8].tolist(), viz[0, :, 0, 0].tolist()
([255, 0, 0], [0, 0, 0])
```

Parameters

 * `canvas` (`Tensor`): `uint8` images of shape `[B, 3, H, W]`. The method does not modify it.
 * `targets` (`Tensor`): Keypoints of shape `[N, 1 + 3 * K]` with rows `[batch_index, x_1, y_1, v_1, ..., v_K]`. The coordinates
   are normalized to `[0, 1]`.
 * `draw_indices` (`bool`): Whether to write the index of each keypoint next to it with
   [draw_keypoint_indices_pil](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md).
   The method fails with `RuntimeError` when an image has more than one instance and `K` is more than `1`.
 * `colors` (`Color`): Color of the keypoints and the indices.
 * `**kwargs`: Keyword arguments for `torchvision.utils.draw_keypoints`, such as `radius` and `connectivity`.

Returns

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

##### forward

```python
def forward(prediction_canvas: Tensor, target_canvas: Tensor, keypoints: list[Tensor], boundingbox: list[Tensor], target_keypoints: Tensor | None, target_boundingbox: Tensor | None, **kwargs) -> tuple[Tensor, Tensor] | Tensor:
```

Draw the predicted boxes and keypoints, and the targets when given.

[BBoxVisualizer.draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/bbox_visualizer.md)
draws the boxes, and
[draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md)
draws the keypoints on top. When `target_boundingbox` is set,
[BBoxVisualizer.draw_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/bbox_visualizer.md)
draws the target boxes. When `target_keypoints` is set,
[draw_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md)
draws the target keypoints on top, in `visible_color`.

When `radius` is `None`, the radius comes from the size of each canvas. It is `1` when both sides are below `96` pixels, `5` when
a side is above `512` pixels, and `2` otherwise.

> **Example**
> ```pycon
>>> import torch
>>> visualizer = KeypointVisualizer(
...     labels=["person"], colors=["red"]
... )
>>> canvas = torch.zeros(1, 3, 16, 16, dtype=torch.uint8)
>>> boxes = [torch.tensor([[2.0, 2.0, 14.0, 14.0, 0.9, 0.0]])]
>>> keypoints = [torch.tensor([[[8.0, 8.0, 0.9]]])]
>>> visualizer(canvas, canvas, keypoints, boxes, None, None).shape
torch.Size([1, 3, 16, 16])
>>> targets = torch.tensor([[0, 0.5, 0.5, 2.0]])
>>> len(
...     visualizer(canvas, canvas, keypoints, boxes, targets, None)
... )
2
```

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.
 * `keypoints` (`list[Tensor]`): One tensor per image, of shape `[M_i, K, 3]`. Each keypoint is `(x, y, confidence)`, with `x` and `y` in pixels. [draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md) scales the coordinates by the `scale` factor.
 * `boundingbox` (`list[Tensor]`): One tensor per image, of shape `[M_i, 6]` with rows `[x1, y1, x2, y2, conf, class]` in pixels. [BBoxVisualizer.draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/bbox_visualizer.md) scales them by the `scale` factor.
 * `target_keypoints` (`Tensor | None`): Keypoints of shape `[N, 1 + 3 * K]` with rows `[batch_index, x_1, y_1, v_1, ..., v_K]`. The coordinates are normalized to `[0, 1]`. `None` when the batch has no `keypoints` labels.
 * `target_boundingbox` (`Tensor | None`): Boxes of shape `[N, 6]` with rows `[batch_index, class, x, y, w, h]`, `xywh` normalized to `[0, 1]`. `None` when the batch has no `boundingbox` labels.
 * `**kwargs`: Keyword arguments for `torchvision.utils.draw_keypoints`, such as `width`. The method passes them to [draw_predictions](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md) and to [draw_targets](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/keypoint_visualizer.md). A key that the method also passes by name, such as `radius` or `connectivity`, raises `TypeError`. [BaseVisualizer.run](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/attached_modules/visualizers/base_visualizer.md) does not pass any.

Returns

 * `tuple[Tensor, Tensor] | Tensor`: The pair `(targets, predictions)` of drawn images when `target_keypoints` or `target_boundingbox` is set; otherwise only the predictions image.

#### Attributes

##### supported_tasks