import collections


def get_label_counts(labels):
  """Creates a JSON compatible dictionary of image per class counts."""
  # JSON does not support integer keys.
  counts = {str(k): v for k, v in collections.Counter(labels.numpy()).items()}
  return counts