import torch


def fill_segmentation_map(mask, cmap):
  """
  Fills a mask of class labels according to the given color map.
  Output is of the same type as the color map.

  :param mask: segmentation mask. Must be of type torch.long and of
      shape NxDxHxW
  :param mapping: must be num_classes x c
  :param ids: when given, colors only these classes; all entries must be < num_classes
  :return: filled segmentation map, shape NxDxHxWxc
  """
  shape = mask.shape
  d = cmap.shape[1]
  indices = mask.long()[..., None].expand(*shape, d).view(mask.numel(), d)
  return torch.gather(cmap, 0, indices).view(*shape, d).to(cmap.dtype)