Merged commit includes the following changes:

-- 434613307 by Waymo Research: Make explanation of joint occupancy-flow metrics more accurate. -- 438083067 by Waymo Research: Change task config for Occupancy Flow challenge. Set `cumulative_waypoints = false`. Note that this changes the challenge parameters. Participants should train models and submit predictions with this setting. Add gamma parameter for amplifying predicted occupancy probabilities when visualizing occupancy grids. -- 438179325 by Waymo Research: Minor fixes for the 3d semseg colab. --
BrunoComCue · Apr 1, 2022 · acd2caf · acd2caf
1 parent eccb0d4
commit acd2caf
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 49 deletions.
diff --git a/tutorial/tutorial_3d_semseg.ipynb b/tutorial/tutorial_3d_semseg.ipynb
@@ -109,7 +109,8 @@
         "id": "ibor0U9XBlX6"
       },
       "source": [
-        "# Read 3D semantic segmentation labels from Frame proto"
+        "# Read 3D semantic segmentation labels from Frame proto\n",
+        " Note that only a subset of the frames have 3d semseg labels."
       ]
     },
     {
@@ -268,7 +269,8 @@
         "      sl_tensor = tf.reshape(tf.convert_to_tensor(sl.data), sl.shape.dims)\n",
         "      sl_points_tensor = tf.gather_nd(sl_tensor, tf.where(range_image_mask))\n",
         "    else:\n",
-        "      sl_points_tensor = tf.zeros([points_tensor.shape[0], 2], dtype=tf.int32)\n",
+        "      num_valid_point = tf.math.reduce_sum(tf.cast(range_image_mask, tf.int32))\n",
+        "      sl_points_tensor = tf.zeros([num_valid_point, 2], dtype=tf.int32)\n",
         "      \n",
         "    point_labels.append(sl_points_tensor.numpy())\n",
         "  return point_labels"
@@ -284,7 +286,7 @@
       "source": [
         "points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n",
         "    frame, range_images, camera_projections, range_image_top_pose)\n",
-        "points_ri2, cp_points_ri2) = frame_utils.convert_range_image_to_point_cloud(\n",
+        "points_ri2, cp_points_ri2 = frame_utils.convert_range_image_to_point_cloud(\n",
         "    frame, range_images, camera_projections, range_image_top_pose, ri_index=1)"
       ],
       "metadata": {

diff --git a/tutorial/tutorial_occupancy_flow.ipynb b/tutorial/tutorial_occupancy_flow.ipynb
@@ -8,17 +8,15 @@
       "source": [
         "# Waymo Open Dataset Occupancy and Flow Tutorial\n",
         "\n",
-        "- Website: https://waymo.com/open\n",
+        "- Website: https://waymo.com/open/challenges/2022/occupancy-flow-prediction-challenge/\n",
         "- GitHub: https://github.com/waymo-research/waymo-open-dataset\n",
         "\n",
         "This tutorial demonstrates:\n",
         "\n",
         "- How to decode and interpret the data in the WOD Motion dataset.\n",
         "- How to construct ground-truth occupancy grids and flow fields.\n",
         "- How to train a simple model with Tensorflow.\n",
-        "- How to run inference on the test set and store the results into a protobuf for submission.\n",
-        "\n",
-        "Visit the [Waymo Open Dataset Website](https://waymo.com/open) to download the full dataset."
+        "- How to run inference on the test set and store the results into a protobuf for submission."
       ]
     },
     {
@@ -32,40 +30,12 @@
         "To run a Jupyter kernel locally, run:\n",
         "\n",
         "```\n",
-        "$ pip install \"waymo_open_dataset_tf_2_6_0==1.4.3\"\n",
+        "$ pip install \"waymo_open_dataset_tf_2_6_0==1.4.4\"\n",
         "$ pip install \"notebook\u003e=5.3\"\n",
         "$ jupyter notebook\n",
         "```"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "xdEcN6WilcBn"
-      },
-      "outputs": [],
-      "source": [
-        "# Data location. Please edit.\n",
-        "\n",
-        "# A tfrecord containing tf.Example protos as downloaded from the Waymo Open\n",
-        "# Dataset (motion) webpage.\n",
-        "\n",
-        "# Replace this path with your own tfrecords.\n",
-        "DATASET_FOLDER = '/path/to/waymo_open_dataset_motion_v_1_1_0/uncompressed'\n",
-        "\n",
-        "# TFRecord dataset.\n",
-        "TRAIN_FILES = f'{DATASET_FOLDER}/tf_example/training/training_tfexample.tfrecord*'\n",
-        "VAL_FILES = f'{DATASET_FOLDER}/tf_example/validation/validation_tfexample.tfrecord*'\n",
-        "TEST_FILES = f'{DATASET_FOLDER}/tf_example/testing/testing_tfexample.tfrecord*'\n",
-        "\n",
-
-        "\n",
-        "# Text files containing validation and test scenario IDs for this challenge.\n",
-        "VAL_SCENARIO_IDS_FILE = f'{DATASET_FOLDER}/occupancy_flow_challenge/validation_scenario_ids.txt'\n",
-        "TEST_SCENARIO_IDS_FILE = f'{DATASET_FOLDER}/occupancy_flow_challenge/testing_scenario_ids.txt'"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -98,6 +68,45 @@
         "from waymo_open_dataset.utils import occupancy_flow_vis"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4Y-PWAeJpTlL"
+      },
+      "source": [
+        "# Data location\n",
+        "\n",
+        "Visit the [Waymo Open Dataset Website](https://waymo.com/open) to download the full dataset."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xdEcN6WilcBn"
+      },
+      "outputs": [],
+      "source": [
+        "# PLEASE EDIT.\n",
+        "\n",
+        "# A tfrecord containing tf.Example protos as downloaded from the Waymo Open\n",
+        "# Dataset (motion) webpage.\n",
+        "\n",
+        "# Replace this path with your own tfrecords.\n",
+        "DATASET_FOLDER = '/path/to/waymo_open_dataset_motion_v_1_1_0/uncompressed'\n",
+        "\n",
+        "# TFRecord dataset.\n",
+        "TRAIN_FILES = f'{DATASET_FOLDER}/tf_example/training/training_tfexample.tfrecord*'\n",
+        "VAL_FILES = f'{DATASET_FOLDER}/tf_example/validation/validation_tfexample.tfrecord*'\n",
+        "TEST_FILES = f'{DATASET_FOLDER}/tf_example/testing/testing_tfexample.tfrecord*'\n",
+        "\n",
+
+        "\n",
+        "# Text files containing validation and test scenario IDs for this challenge.\n",
+        "VAL_SCENARIO_IDS_FILE = f'{DATASET_FOLDER}/occupancy_flow_challenge/validation_scenario_ids.txt'\n",
+        "TEST_SCENARIO_IDS_FILE = f'{DATASET_FOLDER}/occupancy_flow_challenge/testing_scenario_ids.txt'"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -116,7 +125,7 @@
       "outputs": [],
       "source": [
         "filenames = tf.io.matching_files(TRAIN_FILES)\n",
-        "dataset = tf.data.TFRecordDataset(filenames, compression_type='')\n",
+        "dataset = tf.data.TFRecordDataset(filenames)\n",
         "dataset = dataset.repeat()\n",
         "dataset = dataset.map(occupancy_flow_data.parse_tf_example)\n",
         "dataset = dataset.batch(16)\n",
@@ -149,7 +158,7 @@
         "id": "trAv9YGrvYnc"
       },
       "source": [
-        "# Visualize TF Example sample"
+        "# Visualize TF Example"
       ]
     },
     {
@@ -447,7 +456,7 @@
         "num_past_steps: 10\n",
         "num_future_steps: 80\n",
         "num_waypoints: 8\n",
-        "cumulative_waypoints: true\n",
+        "cumulative_waypoints: false\n",
         "normalize_sdc_yaw: true\n",
         "grid_height_cells: 256\n",
         "grid_width_cells: 256\n",
@@ -555,6 +564,7 @@
         "observed_occupancy_rgb = occupancy_flow_vis.occupancy_rgb_image(\n",
         "    agent_grids=observed_occupancy_grids,\n",
         "    roadgraph_image=vis_grids.roadgraph,\n",
+        "    gamma=1.6,\n",
         ")"
       ]
     },
@@ -581,6 +591,7 @@
         "occluded_occupancy_rgb = occupancy_flow_vis.occupancy_rgb_image(\n",
         "    agent_grids=occluded_occupancy_grids,\n",
         "    roadgraph_image=vis_grids.roadgraph,\n",
+        "    gamma=1.6,\n",
         ")"
       ]
     },
@@ -636,6 +647,7 @@
         "  observed_occupancy_rgb = occupancy_flow_vis.occupancy_rgb_image(\n",
         "      agent_grids=observed_occupancy_grids,\n",
         "      roadgraph_image=vis_grids.roadgraph,\n",
+        "      gamma=1.6,\n",
         "  )\n",
         "  images.append(observed_occupancy_rgb[0])\n",
         "\n",
@@ -658,6 +670,7 @@
         "  occluded_occupancy_rgb = occupancy_flow_vis.occupancy_rgb_image(\n",
         "      agent_grids=occluded_occupancy_grids,\n",
         "      roadgraph_image=vis_grids.roadgraph,\n",
+        "      gamma=1.6,\n",
         "  )\n",
         "  images.append(occluded_occupancy_rgb[0])\n",
         "\n",
@@ -1129,7 +1142,7 @@
         "def _apply_sigmoid_to_occupancy_logits(\n",
         "    pred_waypoint_logits: occupancy_flow_grids.WaypointGrids\n",
         ") -\u003e occupancy_flow_grids.WaypointGrids:\n",
-        "  \"\"\"Converts occupancy logits with probabilities.\"\"\"\n",
+        "  \"\"\"Converts occupancy logits to probabilities.\"\"\"\n",
         "  pred_waypoints = occupancy_flow_grids.WaypointGrids()\n",
         "  pred_waypoints.vehicles.observed_occupancy = [\n",
         "      tf.sigmoid(x) for x in pred_waypoint_logits.vehicles.observed_occupancy\n",
@@ -1190,6 +1203,7 @@
         "  observed_occupancy_rgb = occupancy_flow_vis.occupancy_rgb_image(\n",
         "      agent_grids=observed_occupancy_grids,\n",
         "      roadgraph_image=vis_grids.roadgraph,\n",
+        "      gamma=1.6,\n",
         "  )\n",
         "  images.append(observed_occupancy_rgb[0])\n",
         "\n",
@@ -1221,6 +1235,7 @@
         "  occluded_occupancy_rgb = occupancy_flow_vis.occupancy_rgb_image(\n",
         "      agent_grids=occluded_occupancy_grids,\n",
         "      roadgraph_image=vis_grids.roadgraph,\n",
+        "      gamma=1.6,\n",
         "  )\n",
         "  images.append(observed_occupancy_rgb[0])\n",
         "\n",
@@ -1443,7 +1458,7 @@
       "source": [
         "def _make_test_dataset(test_shard_path: str) -\u003e tf.data.Dataset:\n",
         "  \"\"\"Makes a dataset for one shard in the test set.\"\"\"\n",
-        "  test_dataset = tf.data.TFRecordDataset(test_shard_path, compression_type='')\n",
+        "  test_dataset = tf.data.TFRecordDataset(test_shard_path)\n",
         "  test_dataset = test_dataset.map(occupancy_flow_data.parse_tf_example)\n",
         "  test_dataset = test_dataset.batch(1)\n",
         "  return test_dataset"
@@ -1625,7 +1640,7 @@
       "source": [
         "## Submit\n",
         "\n",
-        "Please see the challenge website for instructions on how to submit."
+        "Please see the [challenge website](https://waymo.com/open/challenges/2022/occupancy-flow-prediction-challenge/) for instructions on how to submit."
       ]
     }
   ],

diff --git a/waymo_open_dataset/protos/occupancy_flow_metrics.proto b/waymo_open_dataset/protos/occupancy_flow_metrics.proto
@@ -52,7 +52,7 @@ message OccupancyFlowTaskConfig {
   //
   // The code provided in occupancy_flow_data.py implements the above logic to
   // construct the ground truth.
-  optional bool cumulative_waypoints = 4 [default = true];
+  optional bool cumulative_waypoints = 4 [default = false];
 
   // -- PREDICTION RESOLUTION AND SCALE --
 
@@ -115,7 +115,7 @@ message OccupancyFlowMetrics {
   // The flow-warped occupancy metrics verify correctness of both predicted
   // flow fields and predicted occupancy grids.
   //
-  // Here is how the flow-warped occupancy metrics are computed for waypoint t:
+  // Here is how the flow-warped occupancy metrics are computed for waypoint k:
   //
   // Let:
   // ~O^b_k denote the predicted occupancy at waypoint k of observed vehicles,
@@ -145,15 +145,20 @@ message OccupancyFlowMetrics {
   // ~W_k = ~F_k o O_{k-1},
   //
   // where o indicates function application -- applying the flow field as a
-  // function to transform the occupancy.  If the predicted flow is accurate, we
-  // should also have ~W_k == O_k.  To combine both expectations, we multiply
-  // ~W_k element-wise with ~O_k, to get
+  // function to transform the occupancy.  If the predicted flow is accurate, it
+  // should be able to reach and cover the future occupancy O_k.  Note that
+  // since we predict backward flow fields, ~W_k may predict expansion of
+  // occupancy in different directions and reach a larger area beyond O_k.
+  // Therefore we multiply ~W_k element-wise with ~O_k, to get
   //
   // ~W_k * ~O_k.
   //
   // If the predicted occupancy and flow at waypoint k are accurate, this term
-  // should be equal to the ground-truth O_k.  Therefore, the flow-warped
-  // metrics compute AUC and Soft-IoU between ~W_k * ~O_k and ground-truth O_k.
+  // should be equal to the ground-truth O_k.  In other words, for a grid cell
+  // to be marked as occupied in ~W_k * ~O_k, it should be supported by both
+  // occupancy and flow predictions.  Therefore, the flow-warped metrics compute
+  // AUC and Soft-IoU between ~W_k * ~O_k and ground-truth O_k.
+
   optional float vehicles_flow_warped_occupancy_auc = 6;
   optional float vehicles_flow_warped_occupancy_iou = 7;
 }
diff --git a/waymo_open_dataset/utils/occupancy_flow_vis.py b/waymo_open_dataset/utils/occupancy_flow_vis.py
@@ -25,13 +25,15 @@
 def occupancy_rgb_image(
     agent_grids: occupancy_flow_data.AgentGrids,
     roadgraph_image: tf.Tensor,
+    gamma: float = 1.6,
 ) -> tf.Tensor:
   """Visualize predictions or ground-truth occupancy.
 
   Args:
     agent_grids: AgentGrids object containing optional
       vehicles/pedestrians/cyclists.
     roadgraph_image: Road graph image [batch_size, height, width, 1] float32.
+    gamma: Amplify predicted probabilities so that they are easier to see.
 
   Returns:
     [batch_size, height, width, 3] float32 RGB image.
@@ -44,6 +46,10 @@ def occupancy_rgb_image(
   ped = zeros if agents.pedestrians is None else agents.pedestrians
   cyc = zeros if agents.cyclists is None else agents.cyclists
 
+  veh = tf.math.pow(veh, 1 / gamma)
+  ped = tf.math.pow(ped, 1 / gamma)
+  cyc = tf.math.pow(cyc, 1 / gamma)
+
   # Convert layers to RGB.
   rg_rgb = tf.concat([zeros, zeros, zeros], axis=-1)
   veh_rgb = tf.concat([veh, zeros, zeros], axis=-1)  # Red.