### Install Waymo Open Dataset Libraries and Jupyter

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

Installs the necessary Waymo Open Dataset TensorFlow package and Jupyter notebook for local execution. Ensure you have pip and notebook installed.

```shell
$ pip install waymo-open-dataset-tf-2-12-0==1.6.7
$ pip install "notebook>=5.3"
$ jupyter notebook
```

--------------------------------

### Install Minimal Deeplab2 for WOD (Shell)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_2d_pvps.ipynb

Installs a minimal version of the deeplab2 library specifically for Waymo Open Dataset (WOD) usage. This is optional and only needed if you don't have deeplab2 installed or require specific WOD modules. It downloads and executes an installation script.

```shell
!wget https://raw.githubusercontent.com/waymo-research/waymo-open-dataset/master/src/waymo_open_dataset/pip_pkg_scripts/install_deeplab2.sh -O - | bash
```

--------------------------------

### Compute Camera Keypoint Metrics (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

This example demonstrates how to compute keypoint metrics for camera data. It defines a function `get_camera_data` to extract keypoints and bounding boxes from a `Frame` proto. It then simulates noisy predicted keypoints and uses `keypoint_metrics.create_combined_metric` to calculate and print various metrics such as OKS-based AP.

```Python
#@title Example how to compute metrics for camera keypoints
from typing import Tuple

def get_camera_data(
    frame: dataset_pb2.Frame
) -> Tuple[keypoint_data.KeypointsTensors, keypoint_data.KeypointsTensors]:
  """Extracts camera keypoints and bounding boxes from the input Frame proto."""
  all_keypoints = []
  all_boxes = []
  for cl in frame.camera_labels:
    for l in cl.labels:
      if l.HasField('camera_keypoints'):
        box = keypoint_data.create_camera_box_tensors(l.box, dtype=tf.float32)
        keypoints = keypoint_data.create_camera_keypoints_tensors(
            l.camera_keypoints.keypoint,
            default_location=box.center,
            order=keypoint_data.CANONICAL_ORDER_CAMERA,
            dtype=tf.float32)
        all_keypoints.append(keypoints)
        all_boxes.append(box)
  keypoint_tensors = keypoint_data.stack_keypoints(all_keypoints)
  box_tensors = keypoint_data.stack_boxes(all_boxes)
  return keypoint_tensors, box_tensors


gt_cam, gt_cam_box = get_camera_data(frame)

noise_stddev = 5.0  # in pixels
pr_cam = keypoint_data.KeypointsTensors(
    location=gt_cam.location +
    tf.random.normal(gt_cam.location.shape, stddev=noise_stddev),
    visibility=gt_cam.visibility)

all_metrics = keypoint_metrics.create_combined_metric(
    keypoint_metrics.DEFAULT_CONFIG_CAMERA)
all_metrics.update_state([gt_cam, pr_cam, gt_cam_box])
result = all_metrics.result()

print('Camera keypoint metrics:')
for name, tensor in sorted(result.items(), key=lambda e: e[0]):
  print(f'{name:20s}: {tensor.numpy():.3f}')
```

--------------------------------

### Install Bazelisk

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This shell command installs Bazelisk globally using npm. Bazelisk is a tool that automatically finds and downloads the correct Bazel version for a project.

```bash
!npm install -g @bazel/bazelisk
```

--------------------------------

### Install Patchelf and Run Bazel Tests

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This sequence of shell commands first installs the `patchelf` utility using pip. It then navigates to the project's source directory and runs Bazel tests for the `waymo_open_dataset` library, with progress rate limit set to 10.0.

```bash
!pip3 install patchelf
!cd /content/waymo-od/src/ && bazelisk test //waymo_open_dataset/... --show_progress_rate_limit=10.0
```

--------------------------------

### Install Waymo Open Dataset and Notebook Dependencies

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_womd_camera.ipynb

Installs the Waymo Open Dataset TensorFlow package, notebook, ipywidgets, and enables the Jupyter notebook HTTP over websockets extension for local execution.

```bash
python3 -m pip install waymo-open-dataset-tf-2-12-0==1.6.7
python3 -m pip install "notebook>=5.3" "ipywidgets>=7.5"
python3 -m pip install --upgrade jupyter_http_over_ws>=0.0.7 && \
jupyter serverextension enable --py jupyter_http_over_ws
jupyter notebook
```

--------------------------------

### Display Example 3D Point Cloud Image

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This snippet displays an example of a rendered 3D point cloud. It uses IPython.display to show an image file. Note that the tool used for rendering these point clouds is not publicly available.

```python
from IPython.display import Image, display
display(Image('/content/waymo-od/tutorial/3d_point_cloud.png'))
```

--------------------------------

### Waymo Open Dataset Imports

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Imports essential modules from the Waymo Open Dataset library, including dataset protocols, label definitions, and utility functions for metrics, keypoints, and data transformations.

```python
#@title Waymo Open Dataset imports
from waymo_open_dataset import dataset_pb2
from waymo_open_dataset import label_pb2
from waymo_open_dataset.metrics.python import keypoint_metrics
from waymo_open_dataset.protos import keypoint_pb2
from waymo_open_dataset.utils import box_utils
from waymo_open_dataset.utils import frame_utils
from waymo_open_dataset.utils import keypoint_data
from waymo_open_dataset.utils import keypoint_draw
from waymo_open_dataset.utils import range_image_utils
from waymo_open_dataset.utils import transform_utils
```

--------------------------------

### Compute Laser Keypoint Metrics (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

This example shows how to compute keypoint metrics for laser sensor data. It includes a `get_laser_data` function to extract keypoints and bounding boxes from laser labels in a `Frame` proto. It then simulates noisy predicted keypoints and calculates metrics using `keypoint_metrics.create_combined_metric` with laser-specific configurations.

```Python
#@title Example how to compute metrics for laser keypoints


def get_laser_data(
    frame: dataset_pb2.Frame
) -> Tuple[keypoint_data.KeypointsTensors, keypoint_data.KeypointsTensors]:
  """Extracts laser keypoints and bounding boxes from the input Frame proto."""
  all_keypoints = []
  all_boxes = []
  for l in frame.laser_labels:
    if l.HasField('laser_keypoints'):
      box = keypoint_data.create_laser_box_tensors(l.box, dtype=tf.float32)
      keypoints = keypoint_data.create_laser_keypoints_tensors(
          l.laser_keypoints.keypoint,
          default_location=box.center,
          order=keypoint_data.CANONICAL_ORDER_LASER,
          dtype=tf.float32)
      all_keypoints.append(keypoints)
      all_boxes.append(box)
  keypoint_tensors = keypoint_data.stack_keypoints(all_keypoints)
  box_tensors = keypoint_data.stack_boxes(all_boxes)
  return keypoint_tensors, box_tensors


gt_cam, gt_cam_box = get_laser_data(frame)

noise_stddev = 0.05  # in meters
pr_cam = keypoint_data.KeypointsTensors(
    location=gt_cam.location +
    tf.random.normal(gt_cam.location.shape, stddev=noise_stddev),
    visibility=gt_cam.visibility)

all_metrics = keypoint_metrics.create_combined_metric(
    keypoint_metrics.DEFAULT_CONFIG_LASER)
all_metrics.update_state([gt_cam, pr_cam, gt_cam_box])
result = all_metrics.result()

print('Laser keypoint metrics:')
for name, tensor in sorted(result.items(), key=lambda e: e[0]):
  print(f'{name:20s}: {tensor.numpy():.3f}')
```

--------------------------------

### Clone Waymo Open Dataset Repository

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_local.ipynb

This snippet shows how to clone the Waymo Open Dataset repository using Git. It requires Git to be installed and accessible from the command line.

```shell
git clone https://github.com/waymo-research/waymo-open-dataset.git waymo-open-dataset-repo
cd waymo-open-dataset-repo
```

--------------------------------

### Install Waymo Open Dataset Package

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_maps.ipynb

Installs the Waymo Open Dataset package, specifically version 1.6.7 for TensorFlow 2.12.0. This is a prerequisite for using the dataset's utilities.

```python
!pip install waymo-open-dataset-tf-2-12-0==1.6.7
```

--------------------------------

### Build Waymo Open Dataset Docker Container

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_local.ipynb

This snippet demonstrates how to build a Docker container for the Waymo Open Dataset with a CPU-based Jupyter kernel. Docker must be installed and running.

```shell
docker build --tag=open_dataset_jupyter -f tutorial/cpu-jupyter.Dockerfile .
```

--------------------------------

### Parse TF Example with TensorFlow

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_motion.ipynb

This snippet shows how to load a TFRecord file and parse a single example using TensorFlow. It assumes a predefined `features_description` dictionary is available for parsing.

```python
dataset = tf.data.TFRecordDataset(FILENAME, compression_type='')
data = next(dataset.as_numpy_iterator())
parsed = tf.io.parse_single_example(data, features_description)
```

--------------------------------

### Clone Waymo Open Dataset Repository

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This snippet clones the Waymo Open Dataset GitHub repository and checks out the master branch. It also upgrades pip. Ensure you have git and Python 3 installed.

```shell
!rm -rf waymo-od > /dev/null
!git clone https://github.com/waymo-research/waymo-open-dataset.git waymo-od
!cd waymo-od && git branch -a
!cd waymo-od && git checkout remotes/origin/master
!pip3 install --upgrade pip
```

--------------------------------

### Initialize TensorFlow and Waymo Open Dataset Utilities (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_local.ipynb

This Python code initializes TensorFlow for eager execution and imports necessary utilities from the Waymo Open Dataset library, including frame utilities and dataset protobuf definitions. TensorFlow must be installed.

```python
import tensorflow as tf
import math
import numpy as np
import itertools

from waymo_open_dataset.utils import frame_utils
from waymo_open_dataset import dataset_pb2 as open_dataset

tf.enable_eager_execution()
```

--------------------------------

### Install Waymo Open Dataset Library (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_3d_semseg.ipynb

Installs the Waymo Open Dataset library for TensorFlow 2.12.0. This is a prerequisite for running the tutorial and accessing the dataset utilities.

```python
!pip3 install waymo-open-dataset-tf-2-12-0==1.6.7
```

--------------------------------

### Install Waymo Open Dataset TensorFlow Package

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

Installs the Waymo Open Dataset TensorFlow package using pip within a Jupyter environment. This command ensures the correct version for TensorFlow 2.12.0 is installed.

```python
%pip install waymo-open-dataset-tf-2-12-0==1.6.7
```

--------------------------------

### Start Jupyter Kernel in Docker Container

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_local.ipynb

This snippet starts a Jupyter kernel inside the previously built Docker container, exposing it on port 8888. Docker must be running.

```shell
docker run -p 8888:8888 open_dataset_jupyter
```

--------------------------------

### Setup Compressed Ray Component and Load Data

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_object_asset.ipynb

Sets up the codec for decompressing ray data and loads compressed ray, camera sensor, and LiDAR box information. It merges these datasets and groups them by laser object ID to prepare for visualization.

```python
#@title Setup ray compressed component

# The codec that is used to decompress the ray compressed component.
codec = read_codec(asset_dir, f'{asset_type}_asset_ray_compressed')

asset_ray_compressed_df = read(asset_dir, f'{asset_type}_asset_ray_compressed')
asset_camera_sensor_df = read(asset_dir, f'{asset_type}_asset_camera_sensor')

# Load additional LiDAR box dimensions to obtain the ray-box intersection
laser_box_df = read(dataset_dir, 'lidar_box')

asset_df = v2.merge(asset_camera_sensor_df, asset_ray_compressed_df)
asset_df = v2.merge(asset_df, laser_box_df)

grouped_asset_df = asset_df.groupby(['key.laser_object_id'])
unique_object_df = grouped_asset_df['key.laser_object_id'].unique().compute()

sel_asset_df = grouped_asset_df.get_group(
    unique_object_df.iat[sel_index][0]).reset_index()
```

--------------------------------

### Create PoseEstimationSubmission Proto - Python

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

This Python script demonstrates how to create a PoseEstimationSubmission proto, which is required for the Waymo Open Dataset Pose Estimation Challenge. It includes helper functions to create keypoint and box protos and a fake detector function that generates random pose estimations for each frame. The output is a binproto file suitable for submission.

```python
#@title Create PoseEstimationSubmission proto
import os
import tqdm
import tensorflow as tf
import dataclasses
from typing import Iterable
import numpy as np

from waymo_open_dataset import dataset_pb2
from waymo_open_dataset.protos import keypoint_pb2
from waymo_open_dataset.protos import box_pb2
from waymo_open_dataset import label_pb2
from waymo_open_dataset.utils import keypoint_data
from waymo_open_dataset.protos import keypoints_submission_pb2

# Input dataset
DATASET_FOLDER = '/waymo_open_dataset_'
VALIDATION_FILES = os.path.join(DATASET_FOLDER, 'validation', '*.tfrecord')

# Where results are going to be saved.
OUTPUT_ROOT_DIRECTORY = '/tmp/waymo_pose_estimation/'
os.makedirs(OUTPUT_ROOT_DIRECTORY, exist_ok=True)

# Iterate over all segments of the dataset and collect predictions.
filenames = tf.io.matching_files(VALIDATION_FILES)


def _create_keypoint_proto(
    loc: np.ndarray, kp_type: keypoint_pb2.KeypointType
) -> keypoint_pb2.LaserKeypoint:
  return keypoint_pb2.LaserKeypoint(
      type=kp_type,
      keypoint_3d={
          'location_m': {'x': loc[0], 'y': loc[1], 'z': loc[2]},
          'visibility': {'is_occluded': False},
      },
  )


def _create_all_keypoints_proto(
    all_loc: np.ndarray, all_vis: np.ndarray
) -> keypoint_pb2.LaserKeypoints:
  keypoints = keypoint_pb2.LaserKeypoints()
  for loc, vis, kp_type in zip(
      all_loc, all_vis, keypoint_data.CANONICAL_ORDER_LASER
  ):
    if vis == 0:
      continue
    keypoints.keypoint.append(_create_keypoint_proto(loc, kp_type))
  return keypoints


def _create_box_proto(
    center: np.ndarray, size: np.ndarray, heading: float
) -> box_pb2.Box3d:
  return box_pb2.Box3d(
      center={'x': center[0], 'y': center[1], 'z': center[2]},
      size={'x': size[0], 'y': size[1], 'z': size[2]},
      heading=heading,
  )


# A "detector" which returns a random number of objects with random keypoints.
def fake_pose_estimation(
    frame: dataset_pb2.Frame,
) -> Iterable[keypoints_submission_pb2.PoseEstimation]:
  # An actual detector would use `frame.laser` and `frame.camera_images`,
  # for a fake detector using `tf.random` is good enough.
  num_objects = tf.random.uniform(shape=(), maxval=200, dtype=tf.int32)
  max_num_keypoints = len(keypoint_data.CANONICAL_ORDER_LASER)
  locations = tf.random.uniform(
      shape=[num_objects, max_num_keypoints, 3], dtype=tf.float32
  )
  is_visible = tf.random.uniform(
      shape=[num_objects, max_num_keypoints], maxval=2, dtype=tf.int32
  )
  box_center = tf.random.uniform(shape=[num_objects, 3], dtype=tf.float32)
  box_size = tf.random.uniform(
      shape=[num_objects, 3], maxval=[1, 1, 2], dtype=tf.float32
  )
  box_heading = tf.random.uniform(shape=[num_objects], dtype=tf.float32)

  for all_loc, all_vis, center, size, heading in zip(
      locations, is_visible, box_center, box_size, box_heading
  ):
    yield keypoints_submission_pb2.PoseEstimation(
        key={
            'context_name': frame.context.name,
            'frame_timestamp_micros': frame.timestamp_micros,
        },
        box=_create_box_proto(center.numpy(), size.numpy(), heading.numpy()),
        laser_keypoints=_create_all_keypoints_proto(
            all_loc.numpy(), all_vis.numpy()
        ),
    )


submission = keypoints_submission_pb2.PoseEstimationSubmission(
    account_name='user@example.com',
    unique_method_name='Random Object Generator',
    authors=['First Author', 'Second Author'],
    affiliation='A Real Organization',
    description='Uses tf.random.uniform',
    method_link='http://example.com/project/page.html',
)
num_frames = 0
# To make this faster as part of the tutorial, we will only process 10 Frames.
# Obviously, to create a valid submission, all the frames from the corresponding
# subset of the dataset needs to be processed. 
# NOTE: Frames with labeled keypoints that are missing in the submissions

```

--------------------------------

### Import Waymo Open Dataset and TensorFlow Libraries

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

Imports essential libraries for working with the Waymo Open Dataset, including TensorFlow, plotting utilities, and protobuf definitions. This setup is required for data processing and model training.

```python
import os
import pathlib
from typing import Sequence
import uuid
import zlib

from IPython.display import HTML
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_graphics.image.transformer as tfg_transformer

from google.protobuf import text_format
from waymo_open_dataset.protos import occupancy_flow_metrics_pb2
from waymo_open_dataset.protos import occupancy_flow_submission_pb2
from waymo_open_dataset.protos import scenario_pb2
from waymo_open_dataset.utils import occupancy_flow_data
from waymo_open_dataset.utils import occupancy_flow_grids
from waymo_open_dataset.utils import occupancy_flow_metrics
from waymo_open_dataset.utils import occupancy_flow_renderer
from waymo_open_dataset.utils import occupancy_flow_vis
```

--------------------------------

### Display Sample LiDAR Point Cloud Image

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_womd_lidar.ipynb

Displays an example image of a 3D point cloud using IPython.display.Image. Note that the internal tool used for rendering is not publicly available.

```python
from IPython.display import Image, display
display(Image('/content/waymo-od/tutorial/womd_point_cloud.png'))
```

--------------------------------

### Select Object and Camera for Visualization

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Selects a specific object by its ID and a camera name (e.g., FRONT_RIGHT) from the loaded frame. It then prints the number and details of laser and camera keypoints associated with the selected object.

```python
#@title Select object and camera
object_id = 'DQFLdFau_A8kTPOkDxfgJA'
camera_name = dataset_pb2.CameraName.Name.FRONT_RIGHT

camera_image_by_name = {i.name: i.image for i in frame.images}
obj = labels[object_id]
num_laser_points = len(obj.laser.keypoints.keypoint)
num_camera_points = len(obj.camera[camera_name].keypoints.keypoint)

print(f'Object {object_id} has')
print(f'{num_laser_points} laser keypoints '
      f'(short name | location | is_occluded):')
for k in sorted(obj.laser.keypoints.keypoint, key=lambda k: k.type):
  m = k.keypoint_3d.location_m
  location_str = f'({m.x:.2f}, {m.y:.2f}, {m.z:.2f})'
  print(f'{keypoint_draw.point_name(k.type)}	|'
        f' {location_str:25} | {k.keypoint_3d.visibility.is_occluded}')
print(f'\na LaserKeypoint proto example:\n\n{obj.laser.keypoints.keypoint[0]}')

print(f'{num_camera_points} camera keypoints '
      f'(short name |  location | is_occluded):')
for k in sorted(
    obj.camera[camera_name].keypoints.keypoint, key=lambda k: k.type):
  px = k.keypoint_2d.location_px
  location_str = f'({px.x:.0f}, {px.y:.0f})'
  print(f'{keypoint_draw.point_name(k.type)}	'
        f'| {location_str:13} | {k.keypoint_2d.visibility.is_occluded}')
print(f'\na CameraKeypoint proto example:\n\n'
      f'{obj.camera[camera_name].keypoints.keypoint[0]}')
```

--------------------------------

### Show Laser Keypoints

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Placeholder for code that would display laser keypoints. This snippet is intended to be followed by the actual implementation for visualizing 3D laser keypoints.

```python
#@title Show laser keypoints


```

--------------------------------

### Define Frame Proto File Path

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Specifies the file path to a TFRecords file containing Frame protos with human keypoints. This variable is used later for loading the dataset.

```python
# File path to a tfrecods file with Frame protos with human keypoints.
frame_path = 'frame_with_keypoints.tfrecord'
```

--------------------------------

### Display Camera Keypoints

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Renders and displays camera keypoints for a selected object and camera. It decodes the camera image, crops it to the object's bounding box, and draws a wireframe of the keypoints on the cropped image.

```python
#@title Show camera keypoints
image_np = _imdecode(camera_image_by_name[camera_name])
croped_image, cropped_camera_keypoints = keypoint_draw.crop_camera_keypoints(
    image_np,
    obj.camera[camera_name].keypoints.keypoint,
    obj.camera[camera_name].box,
    margin=0.3)
camera_wireframe = keypoint_draw.build_camera_wireframe(
    cropped_camera_keypoints)

keypoint_draw.OCCLUDED_BORDER_WIDTH = 3
_, ax = plt.subplots(frameon=False, figsize=(5, 7))
_imshow(ax, croped_image)
keypoint_draw.draw_camera_wireframe(ax, camera_wireframe)
```

--------------------------------

### Define Dataset File Paths

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_sim_agents.ipynb

Defines the file paths for the Waymo Open Dataset TFRecords. Users need to replace the placeholder path with their actual dataset location. This setup is crucial for loading the motion data.

```python
# Please edit.

# Replace this path with your own tfrecords.
# This tutorial is based on using data in the Scenario proto format directly,
# so choose the correct dataset version.
DATASET_FOLDER = '/waymo_open_dataset_'

TRAIN_FILES = os.path.join(DATASET_FOLDER, 'training.tfrecord*')
VALIDATION_FILES = os.path.join(DATASET_FOLDER, 'validation.tfrecord*')
TEST_FILES = os.path.join(DATASET_FOLDER, 'test.tfrecord*')
```

--------------------------------

### Setup Visualization Functions (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_object_asset.ipynb

Sets up utility functions for image visualization. `apply_color_mask` overlays a color mask onto an image, while `grid_imshow` displays images in a grid layout. These functions are essential for visualizing processed sensor data.

```python
#@title Setup visualization functions
def apply_color_mask(image: np.ndarray,
                     mask: np.ndarray,
                     color: Tuple[int, int, int],
                     alpha: float = 0.5) -> np.ndarray:
  """Applies the given mask to the image."""
  color = np.array(color)[np.newaxis, :]
  bg = image * (1 - alpha) + alpha * color
  output = np.where(mask, bg, image).astype(np.uint8)
  return output


def grid_imshow(h: int, w: int, images: Any) -> None:
  """Displays images in a grid."""
  fig, axes = plt.subplots(h, w)
  fig.set_size_inches(20, 10)
  fig.tight_layout()
  for i, image in enumerate(images):
    ax = axes[i] if len(images) > 0 else axes
    ax.imshow(image)
    ax.axis('off')

```

--------------------------------

### Visualize All Agent Trajectories

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_motion.ipynb

This Python function visualizes agent trajectories (past, current, and future) from a decoded Waymo Open Dataset example. It stacks state information, creates masks, generates a colormap, and concatenates all states and masks for visualization. It returns a sequence of image arrays.

```python
def visualize_all_agents_smooth(
    decoded_example,
    size_pixels=1000,
):
  """Visualizes all agent predicted trajectories in a serie of images."""
  # Args:
  #   decoded_example: Dictionary containing agent info about all modeled agents.
  #   size_pixels: The size in pixels of the output image.
  # 
  # Returns:
  #   T of [H, W, 3] uint8 np.arrays of the drawn matplotlib's figure canvas.
  
  # [num_agents, num_past_steps, 2] float32.
  past_states = tf.stack(
      [decoded_example['state/past/x'], decoded_example['state/past/y']],
      -1).numpy()
  past_states_mask = decoded_example['state/past/valid'].numpy() > 0.0

  # [num_agents, 1, 2] float32.
  current_states = tf.stack(
      [decoded_example['state/current/x'], decoded_example['state/current/y']],
      -1).numpy()
  current_states_mask = decoded_example['state/current/valid'].numpy() > 0.0

  # [num_agents, num_future_steps, 2] float32.
  future_states = tf.stack(
      [decoded_example['state/future/x'], decoded_example['state/future/y']],
      -1).numpy()
  future_states_mask = decoded_example['state/future/valid'].numpy() > 0.0

  # [num_points, 3] float32.
  roadgraph_xyz = decoded_example['roadgraph_samples/xyz'].numpy()

  num_agents, num_past_steps, _ = past_states.shape
  num_future_steps = future_states.shape[1]

  color_map = get_colormap(num_agents)

  # [num_agens, num_past_steps + 1 + num_future_steps, depth] float32.
  all_states = np.concatenate([past_states, current_states, future_states], 1)

  # [num_agens, num_past_steps + 1 + num_future_steps] float32.
  all_states_mask = np.concatenate(

```

--------------------------------

### Create Waymo Open Dataset v2 Component Objects

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_v2.ipynb

This Python code snippet shows how to instantiate CameraImageComponent, LiDARComponent, CameraBoxComponent, and LiDARBoxComponent objects from a dictionary representation, likely loaded from a dataset row. It's a foundational step for accessing and processing sensor data within the Waymo Open Dataset.

```python
import waymo_open_dataset.v2 as v2

# Assuming 'row' is a dictionary loaded from the dataset
camera_image = v2.CameraImageComponent.from_dict(row)
lidar = v2.LiDARComponent.from_dict(row)
camera_box = v2.CameraBoxComponent.from_dict(row)
lidar_box = v2.LiDARBoxComponent.from_dict(row)
```

--------------------------------

### Load TFRecord Dataset for Occupancy Flow Challenge (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

Loads a TFRecord dataset for the Occupancy Flow challenge. It reads scenario IDs from text files, creates a TensorFlow dataset, parses TF examples, repeats the dataset, and batches it for further processing. Dependencies include TensorFlow and the occupancy_flow_data module.

```python
VAL_SCENARIO_IDS_FILE = f'{DATASET_FOLDER}/occupancy_flow_challenge/validation_scenario_ids.txt'
TEST_SCENARIO_IDS_FILE = f'{DATASET_FOLDER}/occupancy_flow_challenge/testing_scenario_ids.txt'

filenames = tf.io.matching_files(TRAIN_FILES)
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.repeat()
dataset = dataset.map(occupancy_flow_data.parse_tf_example)
dataset = dataset.batch(16)
it = iter(dataset)
```

--------------------------------

### Create Waymo Open Dataset Component Objects in Python

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_v2.ipynb

This snippet demonstrates how to create different component objects (CameraHumanKeypoints, CameraBox, CameraImage, LiDARBox, LiDAR) from a given row of data using the Waymo Open Dataset library. It assumes the `v2` module is available and that `row` is a dictionary-like object containing the necessary data.

```python
import waymo_open_dataset.utils as utils

# Assuming 'row' is loaded from a dataset file
# Example: row = next(iter(dataset.IterateAllRepeatedFields()))

camera_hkp = v2.CameraHumanKeypointsComponent.from_dict(row)
camera_box = v2.CameraBoxComponent.from_dict(row)
camera_image = v2.CameraImageComponent.from_dict(row)
lidar_box = v2.LiDARBoxComponent.from_dict(row)
lidar = v2.LiDARComponent.from_dict(row)
```

--------------------------------

### Install Waymo Open Dataset Library for Google Colab

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Installs the Waymo Open Dataset library with TensorFlow 2.12.0 using pip, suitable for use within a Google Colab environment.

```python
# To run in a colab:
!pip3 install waymo-open-dataset-tf-2-12-0==1.6.7
```

--------------------------------

### Install Waymo Open Dataset Dependencies

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_v2.ipynb

Installs necessary Python packages for working with the Waymo Open Dataset v2, including gcsfs, the dataset package itself, notebook, ipywidgets, and jupyter_http_over_ws. It also enables the jupyter_http_over_ws server extension.

```bash
python3 -m pip install gcsfs waymo-open-dataset-tf-2-12-0==1.6.7
python3 -m pip install "notebook>=5.3" "ipywidgets>=7.5"
python3 -m pip install --upgrade jupyter_http_over_ws>=0.0.7 && \
jupyter serverextension enable --py jupyter_http_over_ws
jupyter notebook
```

--------------------------------

### Prepare Paths for Detection Metrics Computation (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_local.ipynb

This Python code prepares file paths for computing detection metrics using the Waymo Open Dataset's command-line tools. It checks for an existing Jupyter notebook environment and defines paths to prediction, ground truth, binary, and frame files.

```python
# copybara removed file resource import
import os

if os.path.exists('tutorial_local.ipynb'):
    # in case it is executed as a Jupyter notebook from the tutorial folder.
    os.chdir('../')


fake_predictions_path = '{pyglib_resource}waymo_open_dataset/metrics/tools/fake_predictions.bin'.format(pyglib_resource='')
fake_ground_truths_path = '{pyglib_resource}waymo_open_dataset/metrics/tools/fake_ground_truths.bin'.format(pyglib_resource='')
bin_path = '{pyglib_resource}waymo_open_dataset/metrics/tools/compute_detection_metrics_main'.format(pyglib_resource='')
frames_path = '{pyglib_resource}tutorial/frames'.format(pyglib_resource='')
point_cloud_path = '{pyglib_resource}tutorial/3d_point_cloud.png'.format(pyglib_resource='')
```

--------------------------------

### Compute Scenario Metrics for Bundle

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_sim_agents.ipynb

This Python code demonstrates how to load a challenge configuration and compute scenario metrics for a given scenario and its rollouts. It utilizes the metrics module from the Waymo Open Dataset toolkit. The output includes per-feature scores and a final 'metametric' score.

```python
# Load the test configuration.
config = metrics.load_metrics_config(challenge_type)

scenario_metrics = metrics.compute_scenario_metrics_for_bundle(
    config, scenario, scenario_rollouts
)
print(scenario_metrics)
```

--------------------------------

### Install Waymo Open Dataset and Pillow

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

Installs the Waymo Open Dataset TensorFlow package (version 1.6.7 for TF 2.12.0) and Pillow (version 9.2.0). These are necessary for processing Waymo dataset frames.

```shell
!pip3 install waymo-open-dataset-tf-2-12-0==1.6.7
!pip3 install Pillow==9.2.0
```

--------------------------------

### Download Dataset File

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_maps.ipynb

Downloads a sample TFRecord file named 'frames_with_maps.tfrecord' from a specified URL. This file contains frame data with map information necessary for the tutorial.

```shell
!wget https://github.com/waymo-research/waymo-open-dataset/raw/master/tutorial/frames_with_maps.tfrecord
```

--------------------------------

### List All Columns and Types in Waymo Dataset Components (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_v2.ipynb

This Python script iterates through all defined components in the Waymo Open Dataset, retrieves their schema, and prints the name and Arrow data type for each column. It requires the 'waymo-research' library to be installed and accessible. The output is formatted for readability, showing the data type aligned to the left and the column name to the right.

```python
print('Available components:')
for component, tag in v2.TAG_BY_COMPONENT.items():
  print(f'{tag}: {component.__name__}')
  schema = component.schema()
  for column, arrow_type in zip(schema.names, schema.types):
    print(f'\t{str(arrow_type):40s} {column}')
```

--------------------------------

### Compute Scenario Metrics in Python

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_scenario_gen.ipynb

This code demonstrates how to compute scenario metrics for a given bundle using the Waymo metrics library. It loads a configuration, computes metrics for a scenario and its rollouts, and prints the results. Key dependencies include the 'metrics' module from the Waymo Open Dataset.

```python
# Load the test configuration.
config = metrics.load_metrics_config(challenge_type)

scenario_metrics = metrics.compute_scenario_metrics_for_bundle(
    config, scenario, scenario_rollouts, challenge_type
)
print(scenario_metrics)
```

--------------------------------

### Generate Submission Archive

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_sim_agents.ipynb

This Python code outlines the process of generating submission files for the Waymo Open Dataset challenge. It involves creating an output directory, iterating over shards to generate individual binproto files, and then compressing these files into a tar.gz archive for submission. The file naming convention for binproto files is specified.

```python
# Where results are going to be saved.
OUTPUT_ROOT_DIRECTORY = '/tmp/waymo_sim_agents/'
os.makedirs(OUTPUT_ROOT_DIRECTORY, exist_ok=True)
output_filenames = []

# Iterate over shards. This could be parallelized in any custom way, as the
```

--------------------------------

### Compute Detection Metrics via Command Line (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_local.ipynb

This snippet executes the Waymo Open Dataset's command-line tool to compute detection metrics. It requires the `bin_path`, `fake_predictions_path`, and `fake_ground_truths_path` variables to be set correctly.

```python
!{bin_path} {fake_predictions_path} {fake_ground_truths_path}
```

--------------------------------

### Generate Submission Binproto Files in Python

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_scenario_gen.ipynb

This code snippet illustrates the process of generating submission files for the Waymo challenge. It sets up an output directory and prepares to iterate over data shards to create individual binproto files. The naming convention for these files is crucial for validation. Dependencies include the 'os' module.

```python
# Where results are going to be saved.
OUTPUT_ROOT_DIRECTORY = '/tmp/waymo_scenario_gen/'
os.makedirs(OUTPUT_ROOT_DIRECTORY, exist_ok=True)
output_filenames = []

# Iterate over shards. This could be parallelized in any custom way, as the
```

--------------------------------

### Get Next Batch of Data (Python)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

Retrieves the next batch of input data from an iterator. This is typically used after setting up a TensorFlow dataset iterator for training or inference.

```python
inputs = next(it)
```

--------------------------------

### Build PIP Package Wheel

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This shell command initiates the build process for a Python PIP package wheel file using Bazel. It navigates to the source directory and builds the wheel for the `waymo_open_dataset` TensorFlow 2.12.0 package in optimized mode (`-c opt`).

```bash
!cd /content/waymo-od/src/ && bazelisk build -c opt  //waymo_open_dataset/pip_pkg_scripts:wheel
```

--------------------------------

### Creating and Exporting SimAgentsChallengeSubmission

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_scenario_gen.ipynb

This Python code constructs a `SimAgentsChallengeSubmission` object with processed scenario rollouts and metadata. It then serializes this object into a binary proto format and saves it to a file. The submission includes details like account name, method description, and data usage flags.

```python
# Now that we have 2 `ScenarioRollouts` for this shard, we can package them
# into a `SimAgentsChallengeSubmission`. Remember to populate the metadata
# for each shard.
shard_submission = sim_agents_submission_pb2.SimAgentsChallengeSubmission(
    scenario_rollouts=scenario_rollouts,
    submission_type=sim_agents_submission_pb2.SimAgentsChallengeSubmission.SIM_AGENTS_SUBMISSION,
    account_name='your_account@test.com',
    unique_method_name='scenario_gen_tutorial',
    authors=['test'],
    affiliation='waymo',
    description='Submission from the Scenario Gen tutorial',
    method_link='https://waymo.com/open/',
    # New REQUIRED fields.
    uses_lidar_data=False,
    uses_camera_data=False,
    uses_public_model_pretraining=False,
    num_model_parameters='24',
    acknowledge_complies_with_closed_loop_requirement=True,
)

# Now we can export this message to a binproto, saved to local storage.
output_filename = f'submission.binproto{shard_suffix}'
with open(os.path.join(OUTPUT_ROOT_DIRECTORY, output_filename), 'wb') as f:
  f.write(shard_submission.SerializeToString())
output_filenames.append(output_filename)
```

--------------------------------

### List Built Wheel Files

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This shell command lists the contents of the directory where the PIP package wheel file was built by Bazel. It uses `ls -R -l -h` for a recursive, long, and human-readable listing, showing the details of the generated wheel file.

```bash
!ls -R -l -h /content/waymo-od/src/bazel-bin/waymo_open_dataset/pip_pkg_scripts/waymo_open_dataset_tf_2.12.0*
```

--------------------------------

### Import Waymo Dataset Libraries

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_sim_agents.ipynb

Imports essential libraries for data handling, visualization, and simulation within the Waymo Open Dataset framework. Includes TensorFlow, Matplotlib, and specific Waymo utility modules.

```python
# Imports
import os
import tarfile

# Set matplotlib to jshtml so animations work with colab.
from matplotlib import rc
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tqdm

from waymo_open_dataset.protos import scenario_pb2
from waymo_open_dataset.protos import sim_agents_submission_pb2
from waymo_open_dataset.utils import trajectory_utils
from waymo_open_dataset.utils.sim_agents import submission_specs
from waymo_open_dataset.utils.sim_agents import visualizations
from waymo_open_dataset.wdl_limited.sim_agents_metrics import metric_features
from waymo_open_dataset.wdl_limited.sim_agents_metrics import metrics

rc('animation', html='jshtml')
```

--------------------------------

### Load Frame Proto from TFRecord

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Loads a Frame proto from a TFRecord file using TensorFlow's TFRecordDataset. It iterates through the dataset, parses the first record into a Frame object, and groups object labels.

```python
#@title Load Frame proto
import tensorflow as tf

dataset = tf.data.TFRecordDataset(frame_path, compression_type='')
for data in dataset:
  frame = dataset_pb2.Frame()
  frame.ParseFromString(bytearray(data.numpy()))
  break

labels = keypoint_data.group_object_labels(frame)
print(f'Loaded {len(labels)} objects')
```

--------------------------------

### Get Vehicle Pose

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_vision_based_e2e_driving.ipynb

Retrieves the pose of the ego vehicle from the first image's metadata. This pose is assumed to be an identity matrix if world coordinates have already been converted to vehicle coordinates.

```python
vehicle_pose = data.frame.images[0].pose
```

--------------------------------

### Package Submission to Tarball (Shell)

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_camera_only.ipynb

This command-line snippet shows how to create a compressed tar archive of the submission files. It first changes the directory to '/tmp', creates a tar archive named 'MySubmission.tar', and then compresses it using gzip. The resulting file '/tmp/MySubmission.tar.gz' is ready for upload.

```bash
cd /tmp
tar cvf MySubmission.tar MySubmission
gzip MySubmission.tar
```

--------------------------------

### Compute Detection Metrics

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial.ipynb

This shell command executes a Bazel-compiled C++ tool to compute detection metrics. It takes prediction and ground truth binary files as input. This tool is part of the Waymo Open Dataset's metrics computation library.

```bash
!cd /content/waymo-od/src/ && bazel-bin/waymo_open_dataset/metrics/tools/compute_detection_metrics_main waymo_open_dataset/metrics/tools/fake_predictions.bin  waymo_open_dataset/metrics/tools/fake_ground_truths.bin
```

--------------------------------

### Get Predicted Waypoint Logits from Model Outputs

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

Slices the model's output tensor into occupancy and flow grids for each waypoint. It iterates through the predicted channels, separating them into observed occupancy, occluded occupancy, and flow predictions.

```python
def _get_pred_waypoint_logits(
    model_outputs: tf.Tensor,
) -> occupancy_flow_grids.WaypointGrids:
  """Slices model predictions into occupancy and flow grids."""
  pred_waypoint_logits = occupancy_flow_grids.WaypointGrids()

  # Slice channels into output predictions.
  for k in range(config.num_waypoints):
    index = k * NUM_PRED_CHANNELS
    waypoint_channels = model_outputs[
        :, :, :, index : index + NUM_PRED_CHANNELS
    ]
    pred_observed_occupancy = waypoint_channels[:, :, :, :1]
    pred_occluded_occupancy = waypoint_channels[:, :, :, 1:2]
    pred_flow = waypoint_channels[:, :, :, 2:]
    pred_waypoint_logits.vehicles.observed_occupancy.append(
        pred_observed_occupancy
    )
    pred_waypoint_logits.vehicles.occluded_occupancy.append(
        pred_occluded_occupancy
    )
    pred_waypoint_logits.vehicles.flow.append(pred_flow)

  return pred_waypoint_logits
```

--------------------------------

### Packaging Submission Shards into a Tarball

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_scenario_gen.ipynb

This Python code snippet demonstrates how to package multiple submission binproto files into a single tar.gz archive. It uses the `tarfile` module to create a compressed archive, adding each generated submission shard file into it. This is the final step before submitting the results.

```python
# Once we have created all the shards, we can package them directly into a
# tar.gz archive, ready for submission.
with tarfile.open(
    os.path.join(OUTPUT_ROOT_DIRECTORY, 'submission.tar.gz'), 'w:gz'
) as tar:
  for output_filename in output_filenames:
    tar.add(
        os.path.join(OUTPUT_ROOT_DIRECTORY, output_filename),
        arcname=output_filename,
    )
```

--------------------------------

### Process TFRecord Dataset with Python

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

This snippet processes a TFRecord dataset, iterating through frames, estimating poses, and counting keypoints. It utilizes TensorFlow datasets and tqdm for progress visualization. The input is a list of filenames, and the output is a processed submission object.

```python
some_filesnames = filenames[:2]
print(f'Start processing {len(some_filesnames)} segments:')
for buf in tqdm.tqdm(tf.data.TFRecordDataset(some_filesnames, num_parallel_reads=8)):
  num_frames += 1
  frame = dataset_pb2.Frame.FromString(buf.numpy())
  submission.pose_estimations.extend(fake_pose_estimation(frame))

keypoint_count = sum(
    [len(e.laser_keypoints.keypoint) for e in submission.pose_estimations]
)
print(
    f'Generated a submission with {len(submission.pose_estimations)} objects'
    f' and {keypoint_count} keypoints for {num_frames} frames.'
)
```

--------------------------------

### Auxiliary Imports and Utilities for Visualization

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_keypoints.ipynb

Imports necessary libraries and defines utility functions for image decoding, displaying images on Matplotlib axes, drawing 3D laser points on Plotly figures, and creating Plotly figures with custom axis settings.

```python
#@title Auixiliary imports and utils

import os
import math
import numpy as np
from matplotlib import pylab as plt
import plotly.graph_objects as go
import itertools
import PIL.Image
import io
import dataclasses


def _imdecode(buf: bytes) -> np.ndarray:
  with io.BytesIO(buf) as fd:
    pil = PIL.Image.open(fd)
    return np.array(pil)


def _imshow(ax: plt.Axes, image_np: np.ndarray):
  ax.imshow(image_np)
  ax.axis('off')
  ax.set_autoscale_on(False)


def _draw_laser_points(fig: go.Figure,
                       points: np.ndarray,
                       color: str = 'gray',
                       size: int = 3):
  """Visualizes laser points on a plotly figure."""
  fig.add_trace(
      go.Scatter3d(
          mode='markers',
          x=points[:, 0],
          y=points[:, 1],
          z=points[:, 2],
          marker=dict(color=color, size=size)))


def _create_plotly_figure() -> go.Figure:
  """Creates a plotly figure for 3D visualization."""
  fig = go.Figure()
  axis_settings = dict(
      showgrid=False,
      zeroline=False,
      showline=False,
      showbackground=False,
      showaxeslabels=False,
      showticklabels=False)
  fig.update_layout(
      width=600,
      height=600,
      showlegend=False,
      scene=dict(
          aspectmode='data',  # force xyz has same scale,
          xaxis=axis_settings,
          yaxis=axis_settings,
          zaxis=axis_settings,
      ),
  )
  return fig
```

--------------------------------

### Process and Save All Test Shards

Source: https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb

This Python loop iterates through all provided test shard paths. For each shard, it creates a test dataset, generates predictions, and saves the submission to a file. It also prints a sample scenario prediction after processing the first shard.

```python
for i, test_shard_path in enumerate(test_shard_paths):
  print(f'Creating submission for test shard {test_shard_path}...')
  test_dataset = _make_test_dataset(test_shard_path=test_shard_path)
  submission = _make_submission_proto()
  _generate_predictions_for_one_test_shard(
      submission=submission,
      test_dataset=test_dataset,
      test_scenario_ids=test_scenario_ids,
      shard_message=f'{i + 1} of {len(test_shard_paths)}',
  )
  _save_submission_to_file(
      submission=submission, test_shard_path=test_shard_path
  )

  if i == 0:
    print('Sample scenario prediction:\n')
    print(submission.scenario_predictions[-1])
```