code: def custom_decode_labels(mask, num_images=1, num_classes=20):
      """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.
      """
      n, h, w = mask.shape
  
      # import ipdb; ipdb.set_trace()
      assert n >= num_images, (
          "Batch size %d should be greater or equal than number of images to save %d."
          % (
              n,
              num_images,
          )
      )
  
      hair_mask = torch.where(
          mask == 2, torch.ones_like(mask), torch.zeros_like(mask)
      )
  
      face_mask = torch.where(
          mask == 13, torch.ones_like(mask), torch.zeros_like(mask)
      )
  
      return hair_mask, face_mask
  
file_path: insightface/reconstruction/ostec/external/graphonomy/FaceHairMask/graphonomy_inference.py
imports: import cv2
  import torch
  import timeit
  import numpy
  from PIL import Image
  import torch.nn.functional
  from torch.autograd import Variable
  from torchvision import transforms
  from FaceHairMask import graph
  from FaceHairMask import graphonomy_process
label_desc: Use image node_id 1 for calling a global function with example usage: custom_decode_labels(mask, num_images, num_classes) and returns: hair_mask,  face_mask