tutorial/tutorial_keypoints.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3tW2MkdrIn6F",
   "metadata": {
    "id": "3tW2MkdrIn6F"
   },
   "source": [

    "\n",
    "## Installation\n",
    "\n",
    "To run Jupyter notebook locally:\n",
    "\n",
    "```\n",
    "python3 -m pip install waymo-open-dataset-tf-2-11-0==1.6.1\n",
    "python3 -m pip install \"notebook>=5.3\" \"ipywidgets>=7.5\"\n",
    "python3 -m pip install --upgrade jupyter_http_over_ws>=0.0.7 && \\\n",
    "jupyter serverextension enable --py jupyter_http_over_ws\n",
    "jupyter notebook\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "1GQP0BZlQFlk",
   "metadata": {
    "id": "1GQP0BZlQFlk"
   },
   "outputs": [],
   "source": [
    "# To run in a colab:\n",
    "!pip3 install waymo-open-dataset-tf-2-11-0==1.6.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xdEsPZX9LCVo",
   "metadata": {
    "id": "xdEsPZX9LCVo"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "As_l5B3o8sl6",
   "metadata": {
    "cellView": "code",
    "id": "As_l5B3o8sl6"
   },
   "outputs": [],
   "source": [
    "#@title Waymo Open Dataset imports\n",
    "from waymo_open_dataset import dataset_pb2\n",
    "from waymo_open_dataset import label_pb2\n",
    "from waymo_open_dataset.metrics.python import keypoint_metrics\n",
    "from waymo_open_dataset.protos import keypoint_pb2\n",
    "from waymo_open_dataset.utils import box_utils\n",
    "from waymo_open_dataset.utils import frame_utils\n",
    "from waymo_open_dataset.utils import keypoint_data\n",
    "from waymo_open_dataset.utils import keypoint_draw\n",
    "from waymo_open_dataset.utils import range_image_utils\n",
    "from waymo_open_dataset.utils import transform_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "c_kFiqySorjE",
   "metadata": {
    "id": "c_kFiqySorjE"
   },
   "outputs": [],
   "source": [
    "# File path to a tfrecods file with Frame protos with human keypoints.\n",
    "frame_path = 'frame_with_keypoints.tfrecord'\n",
    "\n",

    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "6f2d89b6",
   "metadata": {
    "cellView": "code",
    "id": "6f2d89b6"
   },
   "outputs": [],
   "source": [
    "#@title Load Frame proto\n",
    "import tensorflow as tf\n",
    "\n",
    "dataset = tf.data.TFRecordDataset(frame_path, compression_type='')\n",
    "for data in dataset:\n",
    "  frame = dataset_pb2.Frame()\n",
    "  frame.ParseFromString(bytearray(data.numpy()))\n",
    "  break\n",
    "\n",
    "labels = keypoint_data.group_object_labels(frame)\n",
    "print(f'Loaded {len(labels)} objects')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "_FJKloajRyMt",
   "metadata": {
    "id": "_FJKloajRyMt"
   },
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "6mlBUrd9-g3k",
   "metadata": {
    "id": "6mlBUrd9-g3k"
   },
   "outputs": [],
   "source": [
    "#@title Auixiliary imports and utils\n",
    "\n",
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "from matplotlib import pylab as plt\n",
    "import plotly.graph_objects as go\n",
    "import itertools\n",
    "import PIL.Image\n",
    "import io\n",
    "import dataclasses\n",
    "\n",
    "\n",
    "def _imdecode(buf: bytes) -> np.ndarray:\n",
    "  with io.BytesIO(buf) as fd:\n",
    "    pil = PIL.Image.open(fd)\n",
    "    return np.array(pil)\n",
    "\n",
    "\n",
    "def _imshow(ax: plt.Axes, image_np: np.ndarray):\n",
    "  ax.imshow(image_np)\n",
    "  ax.axis('off')\n",
    "  ax.set_autoscale_on(False)\n",
    "\n",
    "\n",
    "def _draw_laser_points(fig: go.Figure,\n",
    "                       points: np.ndarray,\n",
    "                       color: str = 'gray',\n",
    "                       size: int = 3):\n",
    "  \"\"\"Visualizes laser points on a plotly figure.\"\"\"\n",
    "  fig.add_trace(\n",
    "      go.Scatter3d(\n",
    "          mode='markers',\n",
    "          x=points[:, 0],\n",
    "          y=points[:, 1],\n",
    "          z=points[:, 2],\n",
    "          marker=dict(color=color, size=size)))\n",
    "\n",
    "\n",
    "def _create_plotly_figure() -> go.Figure:\n",
    "  \"\"\"Creates a plotly figure for 3D visualization.\"\"\"\n",
    "  fig = go.Figure()\n",
    "  axis_settings = dict(\n",
    "      showgrid=False,\n",
    "      zeroline=False,\n",
    "      showline=False,\n",
    "      showbackground=False,\n",
    "      showaxeslabels=False,\n",
    "      showticklabels=False)\n",
    "  fig.update_layout(\n",
    "      width=600,\n",
    "      height=600,\n",
    "      showlegend=False,\n",
    "      scene=dict(\n",
    "          aspectmode='data',  # force xyz has same scale,\n",
    "          xaxis=axis_settings,\n",
    "          yaxis=axis_settings,\n",
    "          zaxis=axis_settings,\n",
    "      ),\n",
    "  )\n",
    "  return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "RiBlI9LK2fKs",
   "metadata": {
    "id": "RiBlI9LK2fKs"
   },
   "outputs": [],
   "source": [
    "#@title Select object and camera\n",
    "object_id = 'DQFLdFau_A8kTPOkDxfgJA'\n",
    "camera_name = dataset_pb2.CameraName.Name.FRONT_RIGHT\n",
    "\n",
    "camera_image_by_name = {i.name: i.image for i in frame.images}\n",
    "obj = labels[object_id]\n",
    "num_laser_points = len(obj.laser.keypoints.keypoint)\n",
    "num_camera_points = len(obj.camera[camera_name].keypoints.keypoint)\n",
    "\n",
    "print(f'Object {object_id} has')\n",
    "print(f'{num_laser_points} laser keypoints '\n",
    "      '(short name | location | is_occluded):')\n",
    "for k in sorted(obj.laser.keypoints.keypoint, key=lambda k: k.type):\n",
    "  m = k.keypoint_3d.location_m\n",
    "  location_str = f'({m.x:.2f}, {m.y:.2f}, {m.z:.2f})'\n",
    "  print(f'{keypoint_draw.point_name(k.type)}\\t|'\n",
    "        f' {location_str:25} | {k.keypoint_3d.visibility.is_occluded}')\n",
    "print(f'\\na LaserKeypoint proto example:\\n\\n{obj.laser.keypoints.keypoint[0]}')\n",
    "\n",
    "print(f'{num_camera_points} camera keypoints '\n",
    "      '(short name |  location | is_occluded):')\n",
    "for k in sorted(\n",
    "    obj.camera[camera_name].keypoints.keypoint, key=lambda k: k.type):\n",
    "  px = k.keypoint_2d.location_px\n",
    "  location_str = f'({px.x:.0f}, {px.y:.0f})'\n",
    "  print(f'{keypoint_draw.point_name(k.type)}\\t'\n",
    "        f'| {location_str:13} | {k.keypoint_2d.visibility.is_occluded}')\n",
    "print(f'\\na CameraKeypoint proto example:\\n\\n'\n",
    "      f'{obj.camera[camera_name].keypoints.keypoint[0]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "Zya-PvdW2fKt",
   "metadata": {
    "id": "Zya-PvdW2fKt"
   },
   "outputs": [],
   "source": [
    "#@title Show camera keypoints\n",
    "image_np = _imdecode(camera_image_by_name[camera_name])\n",
    "croped_image, cropped_camera_keypoints = keypoint_draw.crop_camera_keypoints(\n",
    "    image_np,\n",
    "    obj.camera[camera_name].keypoints.keypoint,\n",
    "    obj.camera[camera_name].box,\n",
    "    margin=0.3)\n",
    "camera_wireframe = keypoint_draw.build_camera_wireframe(\n",
    "    cropped_camera_keypoints)\n",
    "\n",
    "keypoint_draw.OCCLUDED_BORDER_WIDTH = 3\n",
    "_, ax = plt.subplots(frameon=False, figsize=(5, 7))\n",
    "_imshow(ax, croped_image)\n",
    "keypoint_draw.draw_camera_wireframe(ax, camera_wireframe)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "Yc-erhA72fKy",
   "metadata": {
    "id": "Yc-erhA72fKy"
   },
   "outputs": [],
   "source": [
    "#@title Show laser keypoints\n",
    "\n",
    "# Select laser points inside pedestrian's bounding box\n",
    "(range_images, camera_projections, _, range_image_top_pose\n",
    ") = frame_utils.parse_range_image_and_camera_projection(frame)\n",
    "points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n",
    "    frame, range_images, camera_projections, range_image_top_pose)\n",
    "points_all = np.concatenate(points, axis=0)\n",
    "box = box_utils.box_to_tensor(obj.laser.box)[tf.newaxis, :]\n",
    "box_points = points_all[box_utils.is_within_box_3d(points_all, box)[:, 0]]\n",
    "print(f'{box_points.shape[0]} laser points selected.')\n",
    "\n",
    "# Visualize 3D scene\n",
    "laser_wireframe = keypoint_draw.build_laser_wireframe(\n",
    "    obj.laser.keypoints.keypoint)\n",
    "fig = _create_plotly_figure()\n",
    "keypoint_draw.draw_laser_wireframe(fig, laser_wireframe)\n",
    "_draw_laser_points(fig, box_points)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bz4MaFMTRnh_",
   "metadata": {
    "id": "bz4MaFMTRnh_"
   },
   "source": [
    "## Metrics\n",
    "\n",
    "`waymo_open_dataset` package supports the following metrics:\n",
    "\n",
    "- Precision (P) and average precision (AP) for various thresholds of Object Keypoint Similarity (OKS).\n",
    "- Percentage of Correct Keypoints (PCK)\n",
    "- Mean Per Joint Position Error MPJPE (aka MPJE)\n",
    "\n",
    "Please refer to [`waymo_open_dataset/metrics/python/keypoint_metrics.py`](https://github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/metrics/keypoint_metrics.py) for implementation details.\n",
    "\n",
    "NOTE: The OKS metric penalizes keypoint coordinates far outside the ground truth bounding box for samples without ground truth keypoints. For such cases we set the default coordinates to be in the middle of the object's box to avoid the penalty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "W3rreupyhnL_",
   "metadata": {
    "id": "W3rreupyhnL_"
   },
   "outputs": [],
   "source": [
    "#@title Example how to compute metrics for camera keypoints\n",
    "from typing import Tuple\n",
    "\n",
    "def get_camera_data(\n",
    "    frame: dataset_pb2.Frame\n",
    ") -> Tuple[keypoint_data.KeypointsTensors, keypoint_data.KeypointsTensors]:\n",
    "  \"\"\"Extracts camera keypoints and bounding boxes from the input Frame proto.\"\"\"\n",
    "  all_keypoints = []\n",
    "  all_boxes = []\n",
    "  for cl in frame.camera_labels:\n",
    "    for l in cl.labels:\n",
    "      if l.HasField('camera_keypoints'):\n",
    "        box = keypoint_data.create_camera_box_tensors(l.box, dtype=tf.float32)\n",
    "        keypoints = keypoint_data.create_camera_keypoints_tensors(\n",
    "            l.camera_keypoints.keypoint,\n",
    "            default_location=box.center,\n",
    "            order=keypoint_data.CANONICAL_ORDER_CAMERA,\n",
    "            dtype=tf.float32)\n",
    "        all_keypoints.append(keypoints)\n",
    "        all_boxes.append(box)\n",
    "  keypoint_tensors = keypoint_data.stack_keypoints(all_keypoints)\n",
    "  box_tensors = keypoint_data.stack_boxes(all_boxes)\n",
    "  return keypoint_tensors, box_tensors\n",
    "\n",
    "\n",
    "gt_cam, gt_cam_box = get_camera_data(frame)\n",
    "\n",
    "noise_stddev = 5.0  # in pixels\n",
    "pr_cam = keypoint_data.KeypointsTensors(\n",
    "    location=gt_cam.location +\n",
    "    tf.random.normal(gt_cam.location.shape, stddev=noise_stddev),\n",
    "    visibility=gt_cam.visibility)\n",
    "\n",
    "all_metrics = keypoint_metrics.create_combined_metric(\n",
    "    keypoint_metrics.DEFAULT_CONFIG_CAMERA)\n",
    "all_metrics.update_state([gt_cam, pr_cam, gt_cam_box])\n",
    "result = all_metrics.result()\n",
    "\n",
    "print('Camera keypoint metrics:')\n",
    "for name, tensor in sorted(result.items(), key=lambda e: e[0]):\n",
    "  print(f'{name:20s}: {tensor.numpy():.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "ZibXsMaXPRrL",
   "metadata": {
    "id": "ZibXsMaXPRrL"
   },
   "outputs": [],
   "source": [
    "#@title Example how to compute metrics for laser keypoints\n",
    "\n",
    "\n",
    "def get_laser_data(\n",
    "    frame: dataset_pb2.Frame\n",
    ") -> Tuple[keypoint_data.KeypointsTensors, keypoint_data.KeypointsTensors]:\n",
    "  \"\"\"Extracts laser keypoints and bounding boxes from the input Frame proto.\"\"\"\n",
    "  all_keypoints = []\n",
    "  all_boxes = []\n",
    "  for l in frame.laser_labels:\n",
    "    if l.HasField('laser_keypoints'):\n",
    "      box = keypoint_data.create_laser_box_tensors(l.box, dtype=tf.float32)\n",
    "      keypoints = keypoint_data.create_laser_keypoints_tensors(\n",
    "          l.laser_keypoints.keypoint,\n",
    "          default_location=box.center,\n",
    "          order=keypoint_data.CANONICAL_ORDER_LASER,\n",
    "          dtype=tf.float32)\n",
    "      all_keypoints.append(keypoints)\n",
    "      all_boxes.append(box)\n",
    "  keypoint_tensors = keypoint_data.stack_keypoints(all_keypoints)\n",
    "  box_tensors = keypoint_data.stack_boxes(all_boxes)\n",
    "  return keypoint_tensors, box_tensors\n",
    "\n",
    "\n",
    "gt_cam, gt_cam_box = get_laser_data(frame)\n",
    "\n",
    "noise_stddev = 0.05  # in meters\n",
    "pr_cam = keypoint_data.KeypointsTensors(\n",
    "    location=gt_cam.location +\n",
    "    tf.random.normal(gt_cam.location.shape, stddev=noise_stddev),\n",
    "    visibility=gt_cam.visibility)\n",
    "\n",
    "all_metrics = keypoint_metrics.create_combined_metric(\n",
    "    keypoint_metrics.DEFAULT_CONFIG_LASER)\n",
    "all_metrics.update_state([gt_cam, pr_cam, gt_cam_box])\n",
    "result = all_metrics.result()\n",
    "\n",
    "print('Laser keypoint metrics:')\n",
    "for name, tensor in sorted(result.items(), key=lambda e: e[0]):\n",
    "  print(f'{name:20s}: {tensor.numpy():.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "3vxkcm3GqQ2w",
   "metadata": {
    "id": "3vxkcm3GqQ2w"
   },
   "outputs": [],
   "source": [
    "#@title Use individual metrics\n",
    "\n",
    "per_type_scales = [\n",
    "    keypoint_metrics.DEFAULT_PER_TYPE_SCALES[t]\n",
    "    for t in keypoint_data.CANONICAL_ORDER_CAMERA\n",
    "]\n",
    "oks = keypoint_metrics.AveragePrecisionAtOKS(per_type_scales, thresholds=[0.95])\n",
    "oks.update_state([gt_cam, pr_cam, gt_cam_box])\n",
    "oks.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fxlDA4P5AyA7",
   "metadata": {
    "id": "fxlDA4P5AyA7"
   },
   "source": [
    "# Generate a submission\n",
    "The Pose Estimation challenge 2023 expects binproto files with `PoseEstimationSubmission` protos.\n",
    "\n",
    "This sections shows an example how to prepare a submission. Follow the instructions on the [challenge web page](https://waymo.com/open/challenges/2023/pose-estimation) to understand how to submit this tar.gz file to our servers for evaluation.\n",
    "\n",
    "*This section is self sufficient, so you don't need to execute any cells above to use it.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "sJFZtgt_BkvX",
   "metadata": {
    "id": "sJFZtgt_BkvX"
   },
   "outputs": [],
   "source": [
    "#@title Create PoseEstimationSubmission proto\n",
    "import os\n",
    "import tqdm\n",
    "import tensorflow as tf\n",
    "import dataclasses\n",
    "from typing import Iterable\n",
    "import numpy as np\n",
    "\n",
    "from waymo_open_dataset import dataset_pb2\n",
    "from waymo_open_dataset.protos import keypoint_pb2\n",
    "from waymo_open_dataset.protos import box_pb2\n",
    "from waymo_open_dataset import label_pb2\n",
    "from waymo_open_dataset.utils import keypoint_data\n",
    "from waymo_open_dataset.protos import keypoints_submission_pb2\n",
    "\n",
    "# Input dataset\n",
    "DATASET_FOLDER = '/waymo_open_dataset_'\n",

    "VALIDATION_FILES = os.path.join(DATASET_FOLDER, 'validation', '*.tfrecord')\n",
    "\n",
    "# Where results are going to be saved.\n",
    "OUTPUT_ROOT_DIRECTORY = '/tmp/waymo_pose_estimation/'\n",
    "os.makedirs(OUTPUT_ROOT_DIRECTORY, exist_ok=True)\n",
    "\n",
    "# Iterate over all segments of the dataset and collect predictions.\n",
    "filenames = tf.io.matching_files(VALIDATION_FILES)\n",
    "\n",
    "\n",
    "def _create_keypoint_proto(\n",
    "    loc: np.ndarray, kp_type: keypoint_pb2.KeypointType\n",
    ") -> keypoint_pb2.LaserKeypoint:\n",
    "  return keypoint_pb2.LaserKeypoint(\n",
    "      type=kp_type,\n",
    "      keypoint_3d={\n",
    "          'location_m': {'x': loc[0], 'y': loc[1], 'z': loc[2]},\n",
    "          'visibility': {'is_occluded': False},\n",
    "      },\n",
    "  )\n",
    "\n",
    "\n",
    "def _create_all_keypoints_proto(\n",
    "    all_loc: np.ndarray, all_vis: np.ndarray\n",
    ") -> keypoint_pb2.LaserKeypoints:\n",
    "  keypoints = keypoint_pb2.LaserKeypoints()\n",
    "  for loc, vis, kp_type in zip(\n",
    "      all_loc, all_vis, keypoint_data.CANONICAL_ORDER_LASER\n",
    "  ):\n",
    "    if vis == 0:\n",
    "      continue\n",
    "    keypoints.keypoint.append(_create_keypoint_proto(loc, kp_type))\n",
    "  return keypoints\n",
    "\n",
    "\n",
    "def _create_box_proto(\n",
    "    center: np.ndarray, size: np.ndarray, heading: float\n",
    ") -> box_pb2.Box3d:\n",
    "  return box_pb2.Box3d(\n",
    "      center={'x': center[0], 'y': center[1], 'z': center[2]},\n",
    "      size={'x': size[0], 'y': size[1], 'z': size[2]},\n",
    "      heading=heading,\n",
    "  )\n",
    "\n",
    "\n",
    "# A \"detector\" which returns a random number of objects with random keypoints.\n",
    "def fake_pose_estimation(\n",
    "    frame: dataset_pb2.Frame,\n",
    ") -> Iterable[keypoints_submission_pb2.PoseEstimation]:\n",
    "  # An actual detector would use `frame.laser` and `frame.camera_images`,\n",
    "  # for a fake detector using `tf.random` is good enough.\n",
    "  num_objects = tf.random.uniform(shape=(), maxval=200, dtype=tf.int32)\n",
    "  max_num_keypoints = len(keypoint_data.CANONICAL_ORDER_LASER)\n",
    "  locations = tf.random.uniform(\n",
    "      shape=[num_objects, max_num_keypoints, 3], dtype=tf.float32\n",
    "  )\n",
    "  is_visible = tf.random.uniform(\n",
    "      shape=[num_objects, max_num_keypoints], maxval=2, dtype=tf.int32\n",
    "  )\n",
    "  box_center = tf.random.uniform(shape=[num_objects, 3], dtype=tf.float32)\n",
    "  box_size = tf.random.uniform(\n",
    "      shape=[num_objects, 3], maxval=[1, 1, 2], dtype=tf.float32\n",
    "  )\n",
    "  box_heading = tf.random.uniform(shape=[num_objects], dtype=tf.float32)\n",
    "\n",
    "  for all_loc, all_vis, center, size, heading in zip(\n",
    "      locations, is_visible, box_center, box_size, box_heading\n",
    "  ):\n",
    "    yield keypoints_submission_pb2.PoseEstimation(\n",
    "        key={\n",
    "            'context_name': frame.context.name,\n",
    "            'frame_timestamp_micros': frame.timestamp_micros,\n",
    "        },\n",
    "        box=_create_box_proto(center.numpy(), size.numpy(), heading.numpy()),\n",
    "        laser_keypoints=_create_all_keypoints_proto(\n",
    "            all_loc.numpy(), all_vis.numpy()\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "submission = keypoints_submission_pb2.PoseEstimationSubmission(\n",
    "    account_name='user@example.com',\n",
    "    unique_method_name='Random Object Generator',\n",
    "    authors=['First Author', 'Second Author'],\n",
    "    affiliation='A Real Organization',\n",
    "    description='Uses tf.random.uniform',\n",
    "    method_link='http://example.com/project/page.html',\n",
    ")\n",
    "num_frames = 0\n",
    "# To make this faster as part of the tutorial, we will only process 10 Frames.\n",
    "# Obviously, to create a valid submission, all the frames from the corresponding\n",
    "# subset of the dataset needs to be processed. \n",
    "# NOTE: Frames with labeled keypoints that are missing in the submissions\n",
    "# will be considered as false negatives and impact the resulting metric.\n",
    "some_filesnames = filenames[:2]\n",
    "print(f'Start processing {len(some_filesnames)} segments:')\n",
    "for buf in tqdm.tqdm(tf.data.TFRecordDataset(some_filesnames, num_parallel_reads=8)):\n",
    "  num_frames += 1\n",
    "  frame = dataset_pb2.Frame.FromString(buf.numpy())\n",
    "  submission.pose_estimations.extend(fake_pose_estimation(frame))\n",
    "\n",
    "keypoint_count = sum(\n",
    "    [len(e.laser_keypoints.keypoint) for e in submission.pose_estimations]\n",
    ")\n",
    "print(\n",
    "    f'Generated a submission with {len(submission.pose_estimations)} objects'\n",
    "    f' and {keypoint_count} keypoints for {num_frames} frames.'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "rtHdhS9pYbGX",
   "metadata": {
    "id": "rtHdhS9pYbGX"
   },
   "outputs": [],
   "source": [
    "#@title Store compressed submission proto\n",
    "import io\n",
    "import tarfile\n",
    "\n",
    "def _add_tar_file(tar: tarfile.TarFile, filename: str, content: bytes) -> None:\n",
    "  \"\"\"Creates a file inside the archive with specified name and content.\"\"\"\n",
    "  info = tarfile.TarInfo(name=filename)\n",
    "  info.size = len(content)\n",
    "  with io.BytesIO(content) as f:\n",
    "    tar.addfile(info, f)\n",
    "\n",
    "\n",
    "submission_path = os.path.join(OUTPUT_ROOT_DIRECTORY, 'submission.tar.gz')\n",
    "submission_buf = submission.SerializeToString()\n",
    "with tarfile.open(submission_path, 'w:gz') as tar:\n",
    "  _add_tar_file(tar, 'submission.binproto', submission_buf)\n",
    "\n",
    "size_mb = os.stat(submission_path).st_size / (1 << 20)\n",
    "uncompressed_size_mb = len(submission_buf) / (1 << 20)\n",
    "print(\n",
    "    f'Stored {size_mb:.2f} Mb (uncompressed {uncompressed_size_mb:.2f} Mb)'\n",
    "    f' submission into {submission_path}'\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "xdEsPZX9LCVo",
    "_FJKloajRyMt",
    "bz4MaFMTRnh_"
   ],
   "name": "Waymo Open Dataset (Perception): Human Keypoints Tutorial",
   "private_outputs": true,
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}