File size: 4,876 Bytes
57746f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
"""
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_"
)
# Load prediction.
# In Pointcept waymo dataset, we minus 1 to label to ignore UNLABELLED class (0 -> -1)
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)
# generate frame submission
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()
|