code: def _get_category_id_map(annotations_root) -> Dict[str, int]:
      """Gets a map from the TAO category id to a tfds category index.
  
      The tfds category index is the index which a category appears in the
      label list.
  
      Args:
        annotations_root: directory containing the train and validation annotations.
  
      Returns:
        id_map: A dict mapping from TAO category id to tfds category index,
          filtered to contain only categories appearing in the train and val set.
      """
  
      train = json.loads((annotations_root / "train.json").read_text())
      val = json.loads(
          (annotations_root / "validation.json").read_text()
      )
  
      merge_map = _merge_categories_map(train)
      merge_map.update(_merge_categories_map(val))
  
      classes = set()
      classes |= {a["category_id"] for a in train["annotations"]}
      classes |= {a["category_id"] for a in val["annotations"]}
  
      id_map = {}
  
      for tfds_id, lvis_id in enumerate(
          sorted([merge_map.get(x, x) for x in classes])
      ):
          id_map[lvis_id] = tfds_id
  
      return id_map
  
imports: from __future__ import annotations
  import collections
  import json
  import os
  from typing import Any, Dict, Optional, Tuple
  from etils import epath
  import numpy
  import tensorflow_datasets.public_api
label_desc: Use image node_id 4 for calling a global function with example usage: _get_category_id_map(annotations_root) and returns: id_map