|
""" |
|
Script for Creating Waymo Semantic Segmentation Submission |
|
|
|
The Waymo dataset toolkit relies on an old version of Tensorflow |
|
which share a conflicting dependency with the Pointcept environment, |
|
therefore we detach the submission generation from the test process |
|
and the script require the following environment: |
|
|
|
```bash |
|
conda create -n waymo python=3.8 -y |
|
conda activate waymo |
|
pip3 install waymo-open-dataset-tf-2-11-0 |
|
``` |
|
|
|
Author: Xiaoyang Wu ([email protected]) |
|
Please cite our work if the code is helpful to you. |
|
""" |
|
|
|
import os |
|
import tqdm |
|
import argparse |
|
import numpy as np |
|
import zlib |
|
import waymo_open_dataset.dataset_pb2 as open_dataset |
|
from waymo_open_dataset.protos import segmentation_metrics_pb2 |
|
from waymo_open_dataset.protos import segmentation_submission_pb2 |
|
|
|
|
|
def compress_array(array: np.ndarray, is_int32: bool = False): |
|
"""Compress a numpy array to ZLIP compressed serialized MatrixFloat/Int32. |
|
|
|
Args: |
|
array: A numpy array. |
|
is_int32: If true, use MatrixInt32, otherwise use MatrixFloat. |
|
|
|
Returns: |
|
The compressed bytes. |
|
""" |
|
if is_int32: |
|
m = open_dataset.MatrixInt32() |
|
else: |
|
m = open_dataset.MatrixFloat() |
|
m.shape.dims.extend(list(array.shape)) |
|
m.data.extend(array.reshape([-1]).tolist()) |
|
return zlib.compress(m.SerializeToString()) |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument( |
|
"--record_path", |
|
required=True, |
|
help="Path to the prediction result folder of Waymo dataset", |
|
) |
|
parser.add_argument( |
|
"--dataset_path", |
|
required=True, |
|
help="Path to the processed Waymo dataset", |
|
) |
|
parser.add_argument( |
|
"--split", |
|
required=True, |
|
choices=["validation", "testing"], |
|
help="Split of the prediction ([training, validation, testing]).", |
|
) |
|
args = parser.parse_args() |
|
file_list = [file for file in os.listdir(args.record_path) if file.endswith(".npy")] |
|
submission = segmentation_submission_pb2.SemanticSegmentationSubmission() |
|
frames = segmentation_metrics_pb2.SegmentationFrameList() |
|
bar = tqdm.tqdm(file_list) |
|
for file in bar: |
|
bar.set_postfix(file=file) |
|
context_name, frame_timestamp_micros = file.strip("segment-*_pred.npy").split( |
|
"_with_camera_labels_" |
|
) |
|
|
|
|
|
pred = np.load(os.path.join(args.record_path, file)) + 1 |
|
masks = np.load( |
|
os.path.join( |
|
args.dataset_path, |
|
args.split, |
|
f"segment-{context_name}_with_camera_labels", |
|
frame_timestamp_micros, |
|
"mask.npy", |
|
), |
|
allow_pickle=True, |
|
) |
|
offset = np.cumsum([mask.sum() for mask in masks.reshape(-1)]) |
|
pred = np.split(pred[: offset[-1]], offset[:-1]) |
|
pred_ri1 = pred[0] |
|
pred_ri2 = pred[5] |
|
mask_ri1 = np.expand_dims(masks[0, 0], -1) |
|
mask_ri2 = np.expand_dims(masks[1, 0], -1) |
|
range_dummy = np.zeros_like(mask_ri1, dtype=np.int32) |
|
range_pred_ri1 = np.zeros_like(mask_ri1, dtype=np.int32) |
|
range_pred_ri1[mask_ri1] = pred_ri1 |
|
range_pred_ri1 = np.concatenate([range_dummy, range_pred_ri1], axis=-1) |
|
range_pred_ri2 = np.zeros_like(mask_ri2, dtype=np.int32) |
|
range_pred_ri2[mask_ri2] = pred_ri2 |
|
range_pred_ri2 = np.concatenate([range_dummy, range_pred_ri2], axis=-1) |
|
|
|
|
|
segmentation_label = open_dataset.Laser() |
|
segmentation_label.name = open_dataset.LaserName.TOP |
|
segmentation_label.ri_return1.segmentation_label_compressed = compress_array( |
|
range_pred_ri1, is_int32=True |
|
) |
|
segmentation_label.ri_return2.segmentation_label_compressed = compress_array( |
|
range_pred_ri2, is_int32=True |
|
) |
|
frame = segmentation_metrics_pb2.SegmentationFrame() |
|
frame.segmentation_labels.append(segmentation_label) |
|
frame.context_name = context_name |
|
frame.frame_timestamp_micros = int(frame_timestamp_micros) |
|
frames.frames.append(frame) |
|
submission.account_name = "***" |
|
submission.unique_method_name = "***" |
|
submission.authors.append("***") |
|
submission.affiliation = "***" |
|
submission.method_link = "***" |
|
submission.sensor_type = ( |
|
segmentation_submission_pb2.SemanticSegmentationSubmission.LIDAR_ALL |
|
) |
|
submission.number_past_frames_exclude_current = 0 |
|
submission.number_future_frames_exclude_current = 0 |
|
submission.inference_results.CopyFrom(frames) |
|
output_filename = os.path.join(args.record_path, "submission.bin") |
|
f = open(output_filename, "wb") |
|
f.write(submission.SerializeToString()) |
|
f.close() |
|
|
