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)