Skip to content

README.md

Latest commit

 

History

History
316 lines (257 loc) · 13.9 KB

README.md

File metadata and controls

316 lines (257 loc) · 13.9 KB
Raw

Waymo Dataset Creation

This directory contains the two-stage pipeline for creating the processed Waymo Open Dataset, starting from the raw TFRecord files on local disk and producing a feature-rich WebDataset ready for training and evaluation.

[Local: raw .tfrecord files]  ──(to_intermediate.py)──►  [Local: intermediate .tar]  ──(to_webdataset_loftup.py)──►  [Local: WebDataset + LoftUp features]
         conda env: waymo                                       conda env: vernata

Two dataset variants

You need to build two separate versions of the dataset:

Variant Output directory Stage 1 --grid_size Stage 2 --no_image_features Used for
SSL pretraining waymo_loftup_webdataset 0.1 (omit — features extracted) SSL pretraining with LoftUp image feature distillation
Downstream segmentation waymo_webdataset_hr none ✓ Linear probing / supervised fine-tuning (full-resolution geometry)
  • The SSL pretraining variant voxel-downsamples to 0.1 m in Stage 1 and extracts DINOv2/LoftUp image features in Stage 2. Since training uses the same grid size anyway, pre-voxelizing avoids storing redundant points and — more importantly — keeps the expensive LoftUp features compact.
  • The downstream segmentation variant stores the full-resolution point cloud (--grid_size none) because validation metrics are computed on the full cloud: the dataloader voxelizes at load time, but then assigns each voxel's predicted class back to every original point in that voxel. Pre-voxelizing would silently discard points and produce incorrect validation scores. Image features are skipped entirely (--no_image_features), which also makes processing much faster and the shards significantly smaller.

⚠️ Do not use the SSL pretraining dataset for downstream segmentation. Validation requires the full-resolution cloud (the dataloader voxelizes, but maps predictions back to every point per voxel). A pre-voxelized dataset would discard points and produce incorrect metrics.

Optional: When training without image feature distillation you can also create a waymo_webdataset variant by passing --no_image_features with --grid_size 0.1. This produces significantly smaller shards than the LoftUp variant and is faster to load.

⚠️ Two separate conda environments are required. The Waymo Open Dataset SDK depends on TensorFlow 2.11 and Python 3.10, which are incompatible with the rest of the project stack. Stage 1 must therefore run in the dedicated waymo environment, while Stage 2 uses the vernata environment.

Stage Script Conda env Python Key dependencies
1 to_intermediate.py waymo 3.10 tensorflow 2.11, waymo-open-dataset-tf-2-11-0 1.6.1, open3d, webdataset
2 to_webdataset_loftup.py vernata 3.12 torch 2.5, loftup, open3d, webdataset

Stage 1: to_intermediate.py — TFRecord Parsing and Intermediate Packaging

What it does

For each segment (a .tfrecord file in the Waymo Open Dataset), this script:

  1. Lists all .tfrecord files under <source_local_dir>/<split>/ for each requested split (training, validation, testing).
  2. Reads each .tfrecord file.
  3. Parses each frame using the Waymo Open Dataset SDK:
    • Decodes LiDAR range images from both return 1 and return 2 into Cartesian 3D coordinates, retaining per-point polar features (range, intensity, elongation).
    • Decodes the 5 camera images (FRONT, FRONT_LEFT, FRONT_RIGHT, SIDE_LEFT, SIDE_RIGHT) and saves them as PNG.
    • For training / validation: skips frames that lack 3D semantic segmentation labels.
    • For testing: skips frames not present in 3d_semseg_test_set_frames.txt.
  4. Processes each frame's point cloud:
    • Translates all points 1.5 m downward (to center the coordinate frame above the road surface).
    • Estimates surface normals using Open3D (radius 0.5 m, max 30 neighbours), oriented toward the sensor origin.
    • Applies voxel downsampling (--grid_size) cropped to ±80 m (XY) / ±10 m (Z).
    • Computes per-point RGB colors by projecting downsampled points onto the camera images using the Waymo camera projection data.
  5. Saves the following arrays into a per-segment WebDataset .tar shard:
    • coord.npy — (N, 3) downsampled point coordinates
    • normal.npy — (N, 3) estimated surface normals
    • strength.npy — (N, 1) tanh-compressed LiDAR intensity
    • color.npy — (M, 3) RGB colors for the colored subset of points
    • projections.npy — (N, 6) raw Waymo camera projection indices and (u, v) coordinates
    • has_color_index.npy — indices into coord of the colored points
    • downsampled_index.npy — indices from the original cloud kept after downsampling
    • segment.npy — (N,) semantic segmentation labels (training/validation only)
    • pose.npy — 4×4 world-from-vehicle transformation matrix
    • {CAMERA_NAME}.png — the 5 decoded camera images
  6. Creates a success marker (.{segment_name}.success) so that already-processed segments are skipped on re-runs.
  7. Writes an index file (waymo_processed_index.txt) listing every split/segment/timestamp entry.

Default commands

SSL pretraining:

conda activate waymo
export WAYMO_OUTPUT_FOLDER=/path/to/waymo

python to_intermediate.py \
    --source_local_dir "/path/to/waymo_tfrecords" \
    --output_dir "$WAYMO_OUTPUT_FOLDER" \
    --splits training validation \
    --num_workers 32 \
    --grid_size 0.1 \
    --verbose

Downstream segmentation (full-resolution cloud — required for correct validation):

conda activate waymo
export WAYMO_OUTPUT_FOLDER=/path/to/waymo

python to_intermediate.py \
    --source_local_dir "/path/to/waymo_tfrecords" \
    --output_dir "${WAYMO_OUTPUT_FOLDER}_no_grid" \
    --splits training validation \
    --num_workers 32 \
    --grid_size none \
    --verbose

Arguments

Argument Default Description
--source_local_dir (required) Local directory containing split sub-directories with .tfrecord files.
--output_dir (required) Local base directory. Intermediate .tar shards and the index file are written to <output_dir>/waymo_intermediate/, with sub-directories training/, validation/, testing/ created automatically.
--index_file waymo_processed_index.txt Filename for the final index listing all processed frame keys.
--splits (required) One or more of training, validation, testing.
--num_workers cpu_count() Number of parallel worker processes (one segment per worker).
--grid_size 0.05 Voxel size (metres) for point cloud downsampling. Pass none to skip downsampling entirely.
--verbose False Enable INFO-level logging.
--partial False If set, process only the first segment per split. Useful for testing the pipeline end-to-end.

Output structure

<output_dir>/
  waymo_intermediate/
    training/
      segment-12345abcde.tar          ← one per TFRecord
      .segment-12345abcde.success     ← marker; segment is skipped on re-run
      ...
    validation/
      ...
    waymo_processed_index.txt         ← lists all "split/segment/timestamp" entries

Each .tar shard contains one sample per Waymo frame, keyed by its timestamp_micros.

Stage 2: to_webdataset_loftup.py — LoftUp Feature Extraction

What it does

Reads the intermediate per-segment .tar files produced by Stage 1 and adds DINOv2/LoftUp visual features to each frame in a single streaming pass:

  1. Scans <source_dir>/<split>/segment-*.tar for each requested split.
  2. Reads each frame's existing data from the intermediate .tar:
    • coord.npy, projections.npy, has_color_index.npy for feature projection.
    • Five camera PNG images (FRONT, FRONT_LEFT, FRONT_RIGHT, SIDE_LEFT, SIDE_RIGHT).
  3. Converts the Waymo (N, 6) projection format into the (5, N, 3) (valid, u, v) format expected by LoftUp's feature projection logic.
  4. Extracts LoftUp features for each frame:
    • Each camera image is resized individually (preserving aspect ratio) to the nearest multiple of the DINOv2 patch size, then passed through the chosen featurizer to produce low-resolution patch features.
    • The LoftUp upsampler network maps these to high-resolution per-pixel features.
    • Features are projected onto 3D points by bilinear sampling at the projected (u, v) pixel coordinates. Points visible in multiple cameras are overwritten in FRONT → FRONT_LEFT → FRONT_RIGHT → SIDE_LEFT → SIDE_RIGHT priority order (earlier cameras take precedence via reverse-iteration overwrite).
  5. Filters the dense (N, D) feature array down to only the colored points (has_color_index) to match the shape of color.npy, producing a sparse (M, D) point_features.npy.
  6. Writes a new .tar for each segment into <output_dir>/<split>/, containing all original data plus point_features.npy.
  7. Creates a success marker so re-runs skip already-finished segments.
  8. Writes an index file (waymo_loftup_index.txt) listing every split/segment/timestamp entry across all processed splits.

Default commands

SSL pretraining (LoftUp image features required):

conda activate vernata
export WAYMO_OUTPUT_FOLDER=/path/to/waymo

python to_webdataset_loftup.py \
    --source_dir "${WAYMO_OUTPUT_FOLDER}/waymo_intermediate" \
    --output_dir "${WAYMO_OUTPUT_FOLDER}/waymo_loftup_webdataset" \
    --splits training validation \
    --num_workers 2 \
    --featurizer_class "dinov2s_reg" \
    --verbose

Downstream segmentation (geometry only, no image features needed):

Image features are not used during downstream segmentation, so skipping LoftUp extraction results in significantly faster processing, smaller output shards, and faster data loading during training. Use the intermediate dataset produced by the downstream Stage 1 command above (with --grid_size none).

conda activate vernata
export WAYMO_OUTPUT_FOLDER=/path/to/waymo

python to_webdataset_loftup.py \
    --source_dir "${WAYMO_OUTPUT_FOLDER}_no_grid/waymo_intermediate" \
    --output_dir "${WAYMO_OUTPUT_FOLDER}_no_grid/waymo_webdataset_hr" \
    --splits training validation \
    --num_workers 8 \
    --no_image_features \
    --verbose

Arguments

Argument Default Description
--source_dir (required) Root directory of the Stage 1 intermediate output (contains training/, validation/ sub-directories).
--output_dir (required) Local directory for output WebDataset shards with LoftUp features.
--splits (required) One or more of training, validation, testing.
--index_file waymo_loftup_index.txt Filename for the final index listing all processed frame keys.
--num_workers cpu_count() Number of parallel worker processes. Set low (e.g., 2) when GPU memory is the bottleneck.
--featurizer_class dinov2s_reg DINOv2 backbone variant to use with LoftUp. Choices: dinov2s, dinov2b, dinov2s_reg, dinov2b_reg.
--no_image_features False Skip LoftUp feature extraction and simply copy each sample as-is. Use this for downstream fine-tuning datasets (see above).
--verbose False Enable INFO-level logging.
--debug False Run in single-process mode (no ProcessPoolExecutor) and load models in the main process. Useful for step-through debugging.

Output structure

<output_dir>/
  training/
    segment-12345abcde.tar          ← original data + point_features.npy
    .segment-12345abcde.success
    ...
  validation/
    ...
  waymo_loftup_index.txt            ← lists all "split/segment/timestamp" entries

Each sample in the output .tar contains all fields from Stage 1 plus:

  • point_features.npy — (M, D) LoftUp feature vectors for the colored points (aligned with color.npy and has_color_index.npy)

Full end-to-end examples

SSL pretraining (with LoftUp features)

export WAYMO_DATA_ROOT=/path/to/waymo_ssl

# ── Stage 1: parse TFRecords and create intermediate dataset ──────────────────
conda activate waymo

python to_intermediate.py \
    --source_local_dir "/path/to/waymo_tfrecords" \
    --output_dir "$WAYMO_DATA_ROOT" \
    --splits training validation \
    --num_workers 32 \
    --grid_size 0.1 \
    --verbose

# ── Stage 2: add LoftUp features ─────────────────────────────────────────────
conda activate vernata

python to_webdataset_loftup.py \
    --source_dir "${WAYMO_DATA_ROOT}/waymo_intermediate" \
    --output_dir "${WAYMO_DATA_ROOT}/waymo_loftup_webdataset" \
    --splits training validation \
    --num_workers 2 \
    --featurizer_class "dinov2s_reg" \
    --verbose

Downstream segmentation (geometry only, full resolution)

export WAYMO_DATA_ROOT=/path/to/waymo_ft

# ── Stage 1: parse TFRecords at full resolution (no voxel downsampling) ───────
conda activate waymo

python to_intermediate.py \
    --source_local_dir "/path/to/waymo_tfrecords" \
    --output_dir "${WAYMO_DATA_ROOT}_no_grid" \
    --splits training validation \
    --num_workers 32 \
    --grid_size none \
    --verbose

# ── Stage 2: copy data without LoftUp features ───────────────────────────────
conda activate vernata

python to_webdataset_loftup.py \
    --source_dir "${WAYMO_DATA_ROOT}/waymo_intermediate" \
    --output_dir "${WAYMO_DATA_ROOT}/waymo_webdataset_hr" \
    --splits training validation \
    --num_workers 8 \
    --no_image_features \
    --verbose