def annotation_3dobject(self, idx, raw=False):
seq_id, frame_idx = idx
fname = "label_lidars/%04d.json" % frame_idx
if self._return_file_path:
return self.base_path / seq_id / fname
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
data = json.loads(ar.read(fname))
else:
with (self.base_path / seq_id / fname).open() as fin:
data = json.load(fin)
labels = list(map(edict, data))
if raw:
return labels
arr = Target3DArray(frame="vehicle") # or frame=None
for label in labels:
tid = base64.urlsafe_b64decode(label.id[:12])
tid, = struct.unpack('Q', tid[:8])
target = ObjectTarget3D(label.center,
Rotation.from_euler("z", label.heading),
label.size,
ObjectTag(label.label, WaymoObjectClass),
tid=tid)
arr.append(target)
return arr
def annotation_3dpoints(self, idx):
seq_id, frame_idx = idx
self._preload_3dsemantics(seq_id)
fnames = {}
for key in self._semantic_dtypes:
fnames[key] = Path("data_3d_semantics", seq_id, "velodyne", key,
"%010d.bin" % frame_idx)
if self._return_file_path:
return edict({k: self.base_path / v for k, v in fnames.items()})
data = edict()
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}_semantics.zip",
to_extract=[str(v)
for v in fnames.values()]) as ar:
for k, v in fnames.items():
data[k] = np.frombuffer(ar.read(str(v)),
dtype=self._semantic_dtypes[k])
else:
for k, v in fnames.items():
data[k] = np.fromfile(self.base_path / v,
dtype=self._semantic_dtypes[k])
data.visible = np.unpackbits(data.visible,
count=len(data.semantic)).astype(bool)
return data
def annotation_3dpoints(self,
idx,
names='lidar_top',
parse_tag=True,
convert_tag=True):
'''
:param parse_tag: Parse tag from original nuscenes id defined in category.json into :class:`NuscenesObjectClass`
:param convert_tag: Convert tag from :class:`NuscenesObjectClass` to :class:`NuscenesSegmentationClass`
'''
assert names == 'lidar_top'
seq_id, frame_idx = idx
fname = "lidar_top_seg/%03d.bin" % frame_idx
if self._return_file_path:
return edict(semantic=self.base_path / seq_id / fname)
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as ar:
buffer = ar.read(fname)
else:
buffer = (self.base_path / seq_id / fname).read_bytes()
label = np.frombuffer(buffer, dtype='u1')
if parse_tag:
if convert_tag:
return edict(semantic=self._segmapping[label])
else:
return edict(semantic=self._rawmapping[label])
else:
return edict(semantic=label)
def _load_camera_data(self, seq_id, fname):
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as ar:
return Image.open(ar.open(fname)).convert('RGB')
else:
return Image.open(self.base_path / seq_id / fname)
def _load_metadata(self):
meta_path = self.base_path / "metadata.msg"
if not meta_path.exists():
_logger.info("Creating metadata of Waymo dataset (%s)...",
self.phase)
metadata = {}
if self.inzip:
for archive in self.base_path.iterdir():
if archive.is_dir() or archive.suffix != ".zip":
continue
with PatchedZipFile(archive,
to_extract="context/stats.json") as ar:
metadata[archive.stem] = json.loads(
ar.read("context/stats.json"))
else:
for folder in self.base_path.iterdir():
if not folder.is_dir():
continue
with (folder / "context/stats.json").open() as fin:
metadata[folder.name] = json.load(fin)
with open(meta_path, "wb") as fout:
msgpack.pack(metadata, fout)
with open(meta_path, "rb") as fin:
self._metadata = SortedDict()
meta_json = msgpack.unpack(fin)
for k, v in meta_json.items():
self._metadata[k] = edict(v)
def timestamp(self, idx):
seq_id, frame_idx = idx
fname = "timestamp/%04d.txt" % frame_idx
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
return int(ar.read(fname).decode())
else:
return int((self.base_path / seq_id / fname).read_bytes())
def timestamp(self, idx, names="lidar_top"):
seq_id, frame_idx = idx
fname = "timestamp/%03d.json" % frame_idx
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as ar:
tsdict = json.loads(ar.read(fname))
else:
with Path(self.base_path, seq_id, fname).open() as fin:
tsdict = json.load(fin)
return tsdict[names]
def _load_lidar_data(self, seq_id, fname):
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as ar:
buffer = ar.read(fname)
else:
buffer = (self.base_path / seq_id / fname).read_bytes()
scan = np.frombuffer(buffer, dtype=np.float32)
scan = np.copy(scan.reshape(-1, 5)) # (x, y, z, intensity, ring index)
return scan
def _load_metadata(self):
meta_path = self.base_path / "metadata.msg"
if not meta_path.exists():
_logger.info("Creating metadata of Nuscenes dataset (%s)...",
self.phase)
metadata = {}
if self.inzip:
for archive in self.base_path.iterdir():
if archive.is_dir() or archive.suffix != ".zip":
continue
with PatchedZipFile(archive,
to_extract="scene/stats.json") as ar:
metadata[archive.stem] = json.loads(
ar.read("scene/stats.json"))
else:
for folder in self.base_path.iterdir():
if not folder.is_dir() or folder.name == "maps":
continue
metadata[folder.name] = json.loads(
(folder / "scene/stats.json").read_text())
assert len(metadata) > 0, "The dataset folder contains no valid frame, "\
"please check path or parameters!"
with open(meta_path, "wb") as fout:
msgpack.pack(metadata, fout)
with open(meta_path, "rb") as fin:
self._metadata = SortedDict()
meta_json = msgpack.unpack(fin)
for k, v in meta_json.items():
self._metadata[k] = edict(v)
# load category mapping for segmentation
with open(self.base_path / "category.json") as fin:
cat_json = json.load(fin)
cat_dict = {}
for item in cat_json:
if 'index' in item:
cat_dict[item['index']] = NuscenesObjectClass.parse(
item['name'])
builtin_table = NuscenesObjectClass._get_nuscenes_id_table()
self._rawmapping = np.empty(len(builtin_table) + 1, dtype='u4')
self._segmapping = np.empty(len(builtin_table) + 1, dtype='u1')
for idx, clsobj in enumerate(builtin_table):
if idx in cat_dict: # test against offcial definition
assert cat_dict[
idx] == clsobj, "Builtin Nuscenes-lidarseg table is incorrect! Please report this bug."
self._rawmapping[idx] = clsobj.value
self._segmapping[idx] = clsobj.to_segmentation().value
def calibration_data(self, idx, raw=False):
uidx = self._parse_idx(idx)
fname = Path(self.phase_path, 'calib', '%06d.txt' % uidx)
if self._return_file_path:
return self.base_path / fname
if self.inzip:
with PatchedZipFile(self.base_path / "data_object_calib.zip",
to_extract=fname) as source:
return self._load_calib(source, uidx, raw)
else:
return self._load_calib(self.base_path, uidx, raw)
def pose(self, idx, raw=False):
seq_id, frame_idx = idx
fname = "pose/%04d.bin" % frame_idx
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
rt = np.frombuffer(ar.read(fname), dtype="f8")
else:
rt = np.fromfile(self.base_path / seq_id / fname, dtype="f8")
if raw:
return rt
return EgoPose(-rt[:3, 3], Rotation.from_matrix(rt[:3, :3]))
def lidar_data(self, idx, names='velo'):
assert names == 'velo'
uidx = self._parse_idx(idx)
fname = Path(self.phase_path, 'velodyne', '%06d.bin' % uidx)
if self._return_file_path:
return self.base_path / fname
if self.inzip:
with PatchedZipFile(self.base_path / "data_object_velodyne.zip",
to_extract=fname) as source:
return utils.load_velo_scan(source, fname)
else:
return utils.load_velo_scan(self.base_path, fname)
def calibration_data(self, idx):
# TODO: add motion compensation, ref: https://github.com/waymo-research/waymo-open-dataset/issues/146
# https://github.com/waymo-research/waymo-open-dataset/issues/79
if isinstance(idx, int):
seq_id, _ = self._locate_frame(idx)
else:
seq_id, _ = idx
assert not self._return_file_path, "The calibration data is not in a single file!"
calib_params = TransformSet("vehicle")
# load json files
fname_cams = "context/calib_cams.json"
fname_lidars = "context/calib_lidars.json"
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=[fname_cams, fname_lidars]) as ar:
calib_cams = json.loads(ar.read(fname_cams))
calib_lidars = json.loads(ar.read(fname_lidars))
else:
with (self.base_path / seq_id / fname_cams).open() as fin:
calib_cams = json.load(fin)
with (self.base_path / seq_id / fname_lidars).open() as fin:
calib_lidars = json.load(fin)
# parse camera calibration
for frame, calib in calib_cams.items():
frame = "camera_" + frame
(fu, fv, cu,
cv), distort = calib['intrinsic'][:4], calib['intrinsic'][4:]
transform = np.array(calib['extrinsic']).reshape(4, 4)
size = (calib['width'], calib['height'])
calib_params.set_intrinsic_pinhole(frame,
size,
cu,
cv,
fu,
fv,
distort_coeffs=distort)
calib_params.set_extrinsic(transform, frame_from=frame)
# parse lidar calibration
for frame, calib in calib_lidars.items():
frame = "lidar_" + frame
calib_params.set_intrinsic_lidar(frame)
transform = np.array(calib['extrinsic']).reshape(4, 4)
calib_params.set_extrinsic(transform, frame_from=frame)
return calib_params
def pose(self, idx, raw=False):
seq_id, frame_idx = idx
date = self._get_date(seq_id)
file_name = Path(date, seq_id, "oxts", "data", "%010d.txt" % frame_idx)
if self._return_file_path:
return self.base_path / file_name
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=file_name) as data:
oxt = utils.load_oxt_file(data, file_name)[0]
else:
oxt = utils.load_oxt_file(self.base_path, file_name)[0]
return utils.parse_pose_from_oxt(oxt)
def lidar_data(self, idx, names='velo'):
seq_id, frame_idx = idx
date = self._get_date(seq_id)
fname = Path(date, seq_id, 'velodyne_points', 'data',
'%010d.bin' % frame_idx)
if self._return_file_path:
return self.base_path / fname
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as source:
return utils.load_velo_scan(source, fname)
else:
return utils.load_velo_scan(self.base_path, fname)
def camera_data(self, idx, names=None):
"""
:param names: frame names of camera to be loaded
"""
seq_id, frame_idx = idx
fname = "%s/%04d.jpg" % (names, frame_idx)
if self._return_file_path:
return self.base_path / seq_id / fname
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
return Image.open(ar.open(fname)).convert('RGB')
else:
Image.open(self.base_path / seq_id / fname).convert('RGB')
def token(self, idx, names='lidar_top'):
'''
Return the sample data token in original Nuscenes data given data index and sensor name
'''
seq_id, frame_idx = idx
fname = "scene/tokens.json"
assert not self._return_file_path, "The tokens are not stored in a single file!"
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
token_data = json.loads(ar.read(fname))
else:
with Path(self.base_path, seq_id, fname).open() as fin:
token_data = json.load(fin)
return token_data[names][frame_idx]
def camera_data(self, idx, names='cam2'):
seq_id, frame_idx = idx
date = self._get_date(seq_id)
fname = Path(date, seq_id, self._frame2folder[names], 'data',
'%010d.png' % frame_idx)
if self._return_file_path:
return self.base_path / fname
gray = names in ['cam0', 'cam1']
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as source:
return utils.load_image(source, fname, gray=gray)
else:
return utils.load_image(self.base_path, fname, gray=gray)
def lidar_data(self, idx, names='velo'):
assert names == 'velo'
seq_id, frame_idx = idx
# load velodyne points
fname = Path(seq_id, "velodyne_points", "data",
'%010d.bin' % frame_idx)
if self._return_file_path:
return self.base_path / "data_3d_raw" / fname
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}_velodyne.zip",
to_extract=fname) as source:
return load_velo_scan(source, fname)
else:
return load_velo_scan(self.base_path / "data_3d_raw", fname)
def camera_data(self, idx, names='cam1'):
seq_id, frame_idx = idx
_, folder_name, dname, _ = self.FRAME_PATH_MAP[names]
fname = Path(seq_id, folder_name, dname, '%010d.png' % frame_idx)
if self._return_file_path:
return self.base_path / "data_2d_raw" / fname
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}_{folder_name}.zip",
to_extract=fname) as source:
return load_image(source, fname, gray=False)
else:
return load_image(self.base_path / "data_2d_raw",
fname,
gray=False)
def _preload_timestamp(self, seq_id):
date = self._get_date(seq_id)
if seq_id in self._timestamp_cache:
return
tsdict = {}
for frame, folder in self._frame2folder.items():
fname = Path(date, seq_id, folder, "timestamps.txt")
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as data:
tsdict[frame] = utils.load_timestamps(
data, fname, formatted=True).astype(int) // 1000
else:
tsdict[frame] = utils.load_timestamps(
self.base_path, fname, formatted=True).astype(int) // 1000
self._timestamp_cache[seq_id] = tsdict
def _preload_timestamps(self, seq, name):
if (seq, name) in self._timestamp_cache:
return
assert seq in self.sequence_ids
folder, subfolder, _, archive = self.FRAME_PATH_MAP[name]
fname = Path(seq, subfolder, "timestamps.txt")
if self.inzip:
with PatchedZipFile(self.base_path / archive,
to_extract=fname) as data:
ts = load_timestamps(data, fname, formatted=True)
else:
ts = load_timestamps(self.base_path / folder,
fname,
formatted=True)
self._timestamp_cache[(seq, name)] = ts.astype(int) // 1000
def _preload_3dobjects(self, seq_id):
assert self.phase in ["training", "validation"
], "Testing set doesn't contains label"
if seq_id in self._3dobjects_mapping:
return
assert seq_id in self.sequence_ids
fname = Path("data_3d_bboxes", "train", f"{seq_id}.xml")
if self.inzip:
with PatchedZipFile(self.base_path / "data_3d_bboxes.zip",
to_extract=fname) as source:
objlist, fmap = load_bboxes(source, fname)
else:
objlist, fmap = load_bboxes(self.base_path, fname)
self._3dobjects_cache[seq_id] = objlist
self._3dobjects_mapping[seq_id] = fmap
def annotation_3dobject(self, idx, raw=False):
assert self.phase_path != "testing", "Testing dataset doesn't contain label data"
uidx = self._parse_idx(idx)
fname = Path(self.phase_path, 'label_2', '%06d.txt' % uidx)
if self._return_file_path:
return self.base_path / fname
if self.inzip:
with PatchedZipFile(self.base_path / "data_object_label_2.zip",
to_extract=fname) as source:
label = load_label(source, fname)
else:
label = load_label(self.base_path, fname)
if raw:
return label
return parse_label(label, self.calibration_data((uidx, ), raw=True))
def annotation_2dobject(self, idx, names=None):
seq_id, frame_idx = idx
fname = "label_%s/%04d.json" % (names, frame_idx)
if self._return_file_path:
return self.base_path / seq_id / fname
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
data = json.loads(ar.read(fname))
else:
with (self.base_path / seq_id / fname).open() as fin:
data = json.load(fin)
labels = list(map(edict, data))
# XXX: currently we don't have a interface for storing 2d object
return labels
def pose(self, idx, names="lidar_top", raw=False):
# Note that here pose always return the pose of the vehicle, names are for different timestamps
seq_id, frame_idx = idx
fname = "pose/%03d.json" % frame_idx
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as ar:
data = json.loads(ar.read(fname))
else:
with Path(self.base_path, seq_id, fname).open() as fin:
data = json.load(fin)
data = data[names]
if raw:
return data
r = Rotation.from_quat(data['rotation'][1:] + [data['rotation'][0]])
t = r.inv().as_matrix().dot(np.array(data['translation']))
return EgoPose(t, r)
def intermediate_data(self, idx, names=None, ninter_frames=None):
seq_id, frame_idx = idx
fname = "intermediate/%03d/meta.json" % frame_idx
# Load meta json
if self.inzip:
with PatchedZipFile(self.base_path / f"{seq_id}.zip",
to_extract=fname) as ar:
meta = json.loads(ar.read(fname))
else:
with Path(self.base_path, seq_id, fname).open() as fin:
meta = json.load(fin)
if not meta:
return []
# Select frames
if ninter_frames is None:
meta = [edict(item) for item in meta[names]]
else:
meta = [edict(item) for item in meta[names][:ninter_frames]]
# parse pose part
for item in meta:
rotation = item.pop("rotation")
rotation = Rotation.from_quat(rotation[1:] + [rotation[0]])
translation = item.pop("translation")
translation = rotation.inv().as_matrix().dot(translation)
item.pose = EgoPose(translation, rotation)
if self._return_file_path:
for item in meta:
item.file = self.base_path / seq_id / "intermediate" / f"{frame_idx:03}" / item.file
return meta
# Load actual data
for item in meta:
fname = "intermediate/%03d/%s" % (frame_idx, item.pop("file"))
if names in self.VALID_CAM_NAMES:
item.data = self._load_camera_data(seq_id, fname)
else: # names in VALID_LIDAR_NAMES:
item.data = self._load_lidar_data(seq_id, fname)
return meta
def calibration_data(self, idx):
seq_id, _ = idx
assert not self._return_file_path, "The calibration is not in a single file!"
calib_params = TransformSet("ego")
calib_fname = "scene/calib.json"
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=calib_fname) as ar:
calib_data = json.loads(ar.read(calib_fname))
else:
with Path(self.base_path, seq_id, calib_fname).open() as fin:
calib_data = json.load(fin)
for frame, calib in calib_data.items():
# set intrinsics
if frame.startswith('cam'):
image_size = (1600, 900
) # currently all the images have the same size
projection = np.array(calib['camera_intrinsic'])
calib_params.set_intrinsic_camera(frame,
projection,
image_size,
rotate=False)
elif frame.startswith('lidar'):
calib_params.set_intrinsic_lidar(frame)
elif frame.startswith('radar'):
calib_params.set_intrinsic_radar(frame)
else:
raise ValueError("Unrecognized frame name.")
# set extrinsics
r = Rotation.from_quat(calib['rotation'][1:] +
[calib['rotation'][0]])
t = np.array(calib['translation'])
extri = np.eye(4)
extri[:3, :3] = r.as_matrix()
extri[:3, 3] = t
calib_params.set_extrinsic(extri, frame_from=frame)
return calib_params
def lidar_data(self, idx, names=None):
# XXX: support return ri2 data
seq_id, frame_idx = idx
fname = "%s/%04d.bin" % (names, frame_idx)
if self._return_file_path:
return self.base_path / seq_id / fname
if self.inzip:
with PatchedZipFile(self.base_path / (seq_id + ".zip"),
to_extract=fname) as ar:
cloud = np.frombuffer(ar.read(fname), dtype='f4')
else:
cloud = np.fromfile(self.base_path / seq_id / fname, dtype='f4')
cloud = cloud.reshape(-1, 5) # x, y, z, intensity, elongation
# point cloud is represented in base frame in Waymo dataset
calib = self.calibration_data(idx)
rt = calib.extrinsics[names]
cloud[:, :3] = cloud[:, :3].dot(rt[:3, :3].T) + rt[:3, 3]
return cloud
def _preload_poses(self, seq):
if seq in self._poses_idx:
return
assert seq in self.sequence_ids
fname = Path("data_poses", seq, "poses.txt")
if self.inzip:
with PatchedZipFile(self.base_path / "data_poses.zip",
to_extract=fname) as data:
plist = np.loadtxt(data.open(str(fname)))
else:
plist = np.loadtxt(self.base_path / fname)
# do interpolation
pose_indices = plist[:, 0].astype(int)
pose_matrices = plist[:, 1:].reshape(-1, 3, 4)
positions = pose_matrices[:, :, 3]
rotations = Rotation.from_matrix(pose_matrices[:, :, :3])
ts_frame = "velo" # the frame used for timestamp extraction
self._preload_timestamps(seq, ts_frame)
timestamps = self._timestamp_cache[(seq, ts_frame)]
fpos = interp1d(timestamps[pose_indices],
positions,
axis=0,
fill_value="extrapolate")
positions = fpos(timestamps)
frot = interp1d(timestamps[pose_indices],
rotations.as_rotvec(),
axis=0,
fill_value="extrapolate")
rotations = frot(timestamps)
self._poses_idx[seq] = set(pose_indices)
self._poses_t[seq] = positions
self._poses_r[seq] = Rotation.from_rotvec(rotations)