class_docstr: A :class:`.Output` which converts the model outputs to the label of the argmax classification per pixel in the
  image for semantic segmentation tasks.
  
  Args:
      labels_map: A dictionary that map the labels ids to pixel intensities.
      visualize: Whether to visualize the image labels.
docstr: Function that given an image with labels ids and their pixel intensity mapping, creates an RGB representation
  for visualisation purposes.
docstr_tok: ["Function","that","given","an","image","with","labels","ids","and","their","pixel","intensity","mapping",",","creates","an","RGB","representation","for","visualisation","purposes","."]
code: def labels_to_image(
      img_labels: Tensor, labels_map: Dict[int, Tuple[int, int, int]]
  ) -> Tensor:
      """Function that given an image with labels ids and their pixel intensity mapping, creates an RGB representation
      for visualisation purposes."""
      assert len(img_labels.shape) == 2, img_labels.shape
      H, W = img_labels.shape
      out = torch.empty(3, H, W, dtype=torch.uint8)
      for label_id, label_val in labels_map.items():
          mask = img_labels == label_id
          for i in range(3):
              out[i].masked_fill_(mask, label_val[i])
      return out
  
imports: import random
  from typing import Any, Dict, Optional, Tuple, Union
  import numpy
  import torch
  from torch import Tensor
  import flash
  from flash.core.data.io.input import DataKeys
  from flash.core.data.io.output import Output
  from flash.core.registry import FlashRegistry
  from flash.core.utilities.imports import _FIFTYONE_AVAILABLE, _MATPLOTLIB_AVAILABLE, _TORCHVISION_AVAILABLE, lazy_import, requires
  from flash.core.utilities.providers import _FIFTYONE
label_desc: Use image node_id 2 for calling the SegmentationLabelsOutput obj's underlying member method code with example usage: obj.labels_to_image(img_labels, labels_map) and returns: out