【Waymo Open Dataset教程——tutorial_2d_pvps】

import immutabledict
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np

from waymo_open_dataset import dataset_pb2 as open_dataset
from waymo_open_dataset.utils import camera_segmentation_utils

FILE_NAME = '/mnt/diskc/data/waymo_v1.4/training/raw_data/segment-1887497421568128425_94_000_114_000_with_camera_labels.tfrecord'

# 读取2D全景分割标签
dataset = tf.data.TFRecordDataset(FILE_NAME, compression_type='')
frames_with_seg = []
sequence_id = None
for ind, data in enumerate(dataset):
    frame = open_dataset.Frame()
    frame.ParseFromString(bytearray(data.numpy()))
    # Save frames which contain CameraSegmentationLabel messages. We assume that
    # if the first image has segmentation labels, all images in this frame will.
    if frame.images[0].camera_segmentation_label.panoptic_label:
        frames_with_seg.append(frame)
        print(ind)
        if sequence_id is None:
            sequence_id = frame.images[0].camera_segmentation_label.sequence_id
        # Collect 3 frames for this demo. However, any number can be used in practice.
        if frame.images[0].camera_segmentation_label.sequence_id != sequence_id or len(frames_with_seg) > 2:
            break
print(len(frames_with_seg))

# Organize the segmentation labels in order from left to right for viz later.
camera_left_to_right_order = [open_dataset.CameraName.SIDE_LEFT,  # 4
                              open_dataset.CameraName.FRONT_LEFT,  # 2
                              open_dataset.CameraName.FRONT,  # 1
                              open_dataset.CameraName.FRONT_RIGHT,  # 3
                              open_dataset.CameraName.SIDE_RIGHT]  # 5
segmentation_protos_ordered = []
for frame in frames_with_seg:
    segmentation_proto_dict = {image.name : image.camera_segmentation_label for image in frame.images}
    segmentation_protos_ordered.append([segmentation_proto_dict[name] for name in camera_left_to_right_order])

# 读取单独全景标签
# Decode a single panoptic label.
panoptic_label_front = camera_segmentation_utils.decode_single_panoptic_label_from_proto(
    segmentation_protos_ordered[0][2]
)

# Separate the panoptic label into semantic and instance labels.
semantic_label_front, instance_label_front = camera_segmentation_utils.decode_semantic_and_instance_labels_from_panoptic_label(
    panoptic_label_front,
    segmentation_protos_ordered[0][2].panoptic_label_divisor
)

 

# 读取摄像头时间一致实例ID全景标签
# The dataset provides tracking for instances between cameras and over time.
# By setting remap_values=True, this function will remap the instance IDs in
# each image so that instances for the same object will have the same ID between
# different cameras and over time.
segmentation_protos_flat = sum(segmentation_protos_ordered, [])
panoptic_labels, is_tracked_masks, panoptic_label_divisor = camera_segmentation_utils.decode_multi_frame_panoptic_labels_from_protos(
    segmentation_protos_flat, remap_values=True
)

# We can further separate the semantic and instance labels from the panoptic
# labels.
NUM_CAMERA_FRAMES = 5
semantic_labels_multiframe = []
instance_labels_multiframe = []
for i in range(0, len(segmentation_protos_flat), NUM_CAMERA_FRAMES):
    semantic_labels = []
    instance_labels = []
    for j in range(NUM_CAMERA_FRAMES):
        semantic_label, instance_label = camera_segmentation_utils.decode_semantic_and_instance_labels_from_panoptic_label(
            panoptic_labels[i + j], panoptic_label_divisor)
        semantic_labels.append(semantic_label)
        instance_labels.append(instance_label)
    semantic_labels_multiframe.append(semantic_labels)
    instance_labels_multiframe.append(instance_labels)

# 可视化全景分割标签
def _pad_to_common_shape(label):
    return np.pad(label, [[1280 - label.shape[0], 0], [0, 0], [0, 0]])

# Pad labels to a common size so that they can be concatenated.
instance_labels = [[_pad_to_common_shape(label) for label in instance_labels] for instance_labels in instance_labels_multiframe]
semantic_labels = [[_pad_to_common_shape(label) for label in semantic_labels] for semantic_labels in semantic_labels_multiframe]
instance_labels = [np.concatenate(label, axis=1) for label in instance_labels]
semantic_labels = [np.concatenate(label, axis=1) for label in semantic_labels]

instance_label_concat = np.concatenate(instance_labels, axis=0)
semantic_label_concat = np.concatenate(semantic_labels, axis=0)
panoptic_label_rgb = camera_segmentation_utils.panoptic_label_to_rgb(
    semantic_label_concat, instance_label_concat)
plt.figure(figsize=(64, 60))
plt.imshow(panoptic_label_rgb)
plt.grid(False)
plt.axis('off')
plt.show()