code: def decode_parsing(
      labels, num_images=1, num_classes=21, is_pred=False
  ):
      """Decode batch of segmentation masks.
  
      Args:
        mask: result of inference after taking argmax.
        num_images: number of images to decode from the batch.
        num_classes: number of classes to predict (including background).
  
      Returns:
        A batch with num_images RGB images of the same size as the input.
      """
      pred_labels = labels[:num_images].clone().cpu().data
      if is_pred:
          pred_labels = torch.argmax(pred_labels, dim=1)
      n, h, w = pred_labels.size()
  
      labels_color = torch.zeros([n, 3, h, w], dtype=torch.uint8)
      for i, c in enumerate(COLORS):
          c0 = labels_color[:, 0, :, :]
          c1 = labels_color[:, 1, :, :]
          c2 = labels_color[:, 2, :, :]
  
          c0[pred_labels == i] = c[0]
          c1[pred_labels == i] = c[1]
          c2[pred_labels == i] = c[2]
  
      return labels_color
  
file_path: insightface/parsing/dml_csr/utils/utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import torchvision
  import torch
  from PIL import Image
  from torch import nn
label_desc: Use image node_id 1 for calling a global function with example usage: decode_parsing(labels, num_images, num_classes, is_pred) and returns: labels_color