# Path: sunds/typing.py



from typing import List, Tuple
from sunds.typing import ArrayDict, TensorLike  # pylint: disable=g-multiple-import
import tensorflow as tf


class MultiSceneBoundaries:
  """Lookup table across scene boundaries."""

  def __init__(self, scene_boundaries: List[ArrayDict]):
    """Constructor.

    Args:
      scene_boundaries: List of scene examples (as numpy arrays).
    """
    # Create the scene_name -> index id mapping
    scene_name2id = {
        ex['scene_name']: i for i, ex in enumerate(scene_boundaries)
    }
    self._table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(
            list(scene_name2id.keys()),
            list(scene_name2id.values()),
        ),
        # Ideally, it would be best if `StaticHashTable` supported not
        # having a default value, and raised a KeyError instead.
        default_value=-1,
    )

    # Stack all scenes values together
    corners = [ex['scene_box'] for ex in scene_boundaries]
    corners = tf.nest.map_structure(lambda *x: tf.stack(x), *corners)
    self._corners = corners


# Path: sunds/typing.py



# Path: sunds/tasks/boundaries_utils.py
from typing import List, Tuple
from sunds.typing import ArrayDict, TensorLike  # pylint: disable=g-multiple-import
import tensorflow as tf



class MultiSceneBoundaries:
  """Lookup table across scene boundaries."""

  def __init__(self, scene_boundaries: List[ArrayDict]):
    """Constructor.

    Args:
      scene_boundaries: List of scene examples (as numpy arrays).
    """
    # Create the scene_name -> index id mapping
    scene_name2id = {
        ex['scene_name']: i for i, ex in enumerate(scene_boundaries)
    }
    self._table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(
            list(scene_name2id.keys()),
            list(scene_name2id.values()),
        ),
        # Ideally, it would be best if `StaticHashTable` supported not
        # having a default value, and raised a KeyError instead.
        default_value=-1,
    )

    # Stack all scenes values together
    corners = [ex['scene_box'] for ex in scene_boundaries]
    corners = tf.nest.map_structure(lambda *x: tf.stack(x), *corners)
    self._corners = corners


  def get_corners(self, scene_name: TensorLike) -> Tuple[tf.Tensor, tf.Tensor]: