diff --git a/dgp/constants.py b/dgp/constants.py index 1ee957ce..7da98edf 100644 --- a/dgp/constants.py +++ b/dgp/constants.py @@ -34,3 +34,18 @@ FEATURE_TYPE_ID_TO_KEY = OrderedDict({v: k for k, v in FEATURE_KEY_TO_TYPE_ID.items()}) # String identifiers for feature types ALL_FEATURE_TYPES = tuple(FEATURE_KEY_TO_TYPE_ID.keys()) + +#EGO AGENT DIMENSIONS +vehicle_name = 'lexus' +vehicle_length = 5.234 +vehicle_width = 1.900 +vehicle_width_with_mirrors = 2.1852 +vehicle_height = 1.68 +vehicle_wheelbase = 3.12 +vehicle_trackwidth = 1.610 +vehicle_wheel_radius = 0.360 +vehicle_fa_to_bumper_dist = 0.950 +vehicle_applanix_origin_to_fa_dist = 3.12 +vehicle_applanix_origin_to_cg_dist = 1.53 +vehicle_applanix_origin_to_r_bumper = 1.164 +vehicle_applanix_origin_height = 0.000 diff --git a/dgp/datasets/prediction_dataset.py b/dgp/datasets/prediction_dataset.py new file mode 100644 index 00000000..2c60f193 --- /dev/null +++ b/dgp/datasets/prediction_dataset.py @@ -0,0 +1,466 @@ +# Copyright 2021 Toyota Research Institute. All rights reserved. +import copy +import itertools +import logging +import os +import time +from collections import defaultdict +from multiprocessing import Pool, cpu_count + +import numpy as np + +from dgp.annotations import ANNOTATION_REGISTRY +from dgp.constants import ANNOTATION_KEY_TO_TYPE_ID +from dgp.datasets.base_dataset import (BaseDataset, DatasetMetadata, SceneContainer) +from dgp.datasets.synchronized_dataset import SynchronizedSceneDataset +from dgp.proto.scene_pb2 import Scene as ScenePb2 +from dgp.utils.protobuf import open_pbobject + + +class ResampledSceneContainer(SceneContainer): + """Object-oriented container for assembling datasets from collections of scenes. + Each scene is resampled from a scene described within a sub-directory with an associated + scene.json file, by a given sampling rate. + """ + def __init__( + self, + scene_path, + directory=None, + autolabeled_scenes=None, + is_datums_synchronized=False, + use_diskcache=True, + sampling_rate=1.0 + ): + """Initialize a scene with a scene object and optionally provide the + directory containing the scene.json to gather additional information + for directory-based dataset loading mode. + + Parameters + ---------- + scene_path: str + Path to the Scene object containing data samples. + + directory: str, default: None + Optional directory containing scene_.json. + + autolabeled_scenes: dict, default: None + Dictionary mapping (defined as:`autolabel_model`/`annotation_key`) to autolabeled SceneContainer. + + is_datums_synchronized: bool, default: False + If True, sample-level synchronization is required i.e. each sample must contain all datums specified in the requested + `datum_names`, and all samples in this scene must contain the same number of datums. + If False, sample-level synchronization is not required i.e. samples are allowed to have different sets of datums. + + use_diskcache: bool, default: True + If True, cache ScenePb2 object using diskcache. If False, save the object in memory. + NOTE: Setting use_diskcache to False would exhaust the memory if have a large number of scenes. + + sampling_rate: float, default: 1.0 + + """ + super().__init__( + scene_path=scene_path, + directory=directory, + autolabeled_scenes=autolabeled_scenes, + is_datums_synchronized=is_datums_synchronized, + use_diskcache=use_diskcache + ) + self.sampling_rate = sampling_rate + + @property + def scene(self): + """ Returns scene. + - If self.use_diskcache is True: returns the cached `_scene` if available, otherwise load the + scene and cache it. + - If self.use_diskcache is False: returns `_scene` in memory if the instance has attribute + `_scene`, otherwise load the scene and save it in memory. + NOTE: Setting use_diskcache to False would exhaust the memory if have a large number of scenes. + """ + if self.use_diskcache: + cached_path = self.scene_path if self.sampling_rate == 1.0 \ + else os.path.join(self.scene_path, f'{self.sampling_rate:.3f}') + if cached_path in SceneContainer.SCENE_CACHE: + _scene = SceneContainer.SCENE_CACHE.get(cached_path) + if _scene is not None: + return _scene + _scene = self.resample_scene(open_pbobject(self.scene_path, ScenePb2)) + SceneContainer.SCENE_CACHE.add(cached_path, _scene) + return _scene + else: + if self._scene is None: + self._scene = self.resample_scene(open_pbobject(self.scene_path, ScenePb2)) + return self._scene + + def resample_scene(self, scene): + """Resample the scene based on given sampling rate. + + Parameters + ---------- + scene: scene_pb2.Scene + scene protobuf data with original sampling rate. + + Returns + ------- + resampled_scene: scene_pb2.Scene + scene protobuf data with giving sampling rate. + """ + resampled_indices = np.linspace(0, len(scene.samples) - 1, int(len(scene.samples) * self.sampling_rate)) + resampled_indices = resampled_indices.astype(np.int32).tolist() + resampled_scene = copy.deepcopy(scene) + resampled_scene.ClearField('samples') + resampled_scene.ClearField('data') + datum_per_sample = len(scene.data) // len(scene.samples) + for r_idx in resampled_indices: + resampled_scene.samples.append(scene.samples[r_idx]) + for d_idx in range(datum_per_sample): + resampled_scene.data.append(scene.data[r_idx * datum_per_sample + d_idx]) + return resampled_scene + + +class PredictionAgentDataset(BaseDataset): + """Dataset for agent-centric prediction use cases, works just like normal SynchronizedSceneDataset, + but guaranteeing trajectory of main agent is present in any fetched sample. + + Parameters + ---------- + scene_dataset_json: str + Full path to the scene dataset json holding collections of paths to scene json. + + split: str, default: 'train' + Split of dataset to read ("train" | "val" | "test" | "train_overfit"). + + datum_names: list, default: None + Select list of datum names for synchronization (see self.select_datums(datum_names)). + + requested_annotations: tuple, default: None + Tuple of annotation types, i.e. ('bounding_box_2d', 'bounding_box_3d'). Should be equivalent + to directory containing annotation from dataset root. + + requested_main_agent_types: tuple, default: 'car' + Tuple of main agent types, i.e. ('car', 'pedestrian'). + The string should be the same as dataset_metadata.ontology_table. + + requested_main_agent_attributes: tuple[str], default: None + Tuple of main agent attributes, i.e. ('moving', 'stopped'). This is predefined per dataset. + By default (None) will include all attributes. + + requested_autolabels: tuple[str], default: None + Currently not supported. + + forward_context: int, default: 0 + Forward context in frames [T+1, ..., T+forward] + + backward_context: int, default: 0 + Backward context in frames [T-backward, ..., T-1] + + min_main_agent_forward: int, default: 0 + Minimum forward samples for main agent. The main-agent will be guaranteed to appear + minimum samples in forward context; i.e., the main-agent will appear in number of + [min_main_agent_forward, forward_context] samples in forward direction. + + min_main_agent_backward: int, default: 0 + Minimum backward samples for main agent. The main-agent will be guaranteed to appear + minimum samples in backward context; i.e., the main-agent will appear in number of + [min_main_agent_backward, backward_context] samples in backward direction. + + generate_depth_from_datum: str, default: None + Datum name of the point cloud. If is not None, then the depth map will be generated for the camera using + the desired point cloud. + + use_3d_trajectories: bool, default: True + Use 3D trajectories (from bounding_box_3d) as main reference of agents. + This requires camera datum with bounding_box_3d annotations. + + batch_per_agent: bool, default: False + Include whole trajectory of an agent in each batch Fetch, this is designed to be used for inference. + If True, backward_context = forward_context = 0 implicitly. + + fps: float, default: -1 + Frame per second during data fetch. -1 means use original fps. + """ + ATTRIBUTE_NAME = 'behavior' + + def __init__( + self, + dataset_json, + split='train', + datum_names=None, + requested_annotations=('bounding_box_3d', ), + requested_main_agent_types=('car', ), + requested_main_agent_attributes=None, + requested_autolabels=None, + forward_context=0, + backward_context=0, + min_main_agent_forward=0, + min_main_agent_backward=0, + generate_depth_from_datum=None, + use_3d_trajectories=True, + batch_per_agent=False, + fps=-1 + ): + self.generate_depth_from_datum = generate_depth_from_datum + self.use_3d_trajectories = use_3d_trajectories + assert len(datum_names) + self.trajectories_reference = 'lidar' if any(['lidar' in datum_name.lower() \ + for datum_name in datum_names]) else datum_names[0].lower() + self.use_3d_trajectories = use_3d_trajectories or self.trajectories_reference == 'lidar' + self.annotation_reference = 'bounding_box_3d' if self.use_3d_trajectories else 'bounding_box_2d' + assert self.annotation_reference in requested_annotations + assert min_main_agent_backward <= backward_context and \ + min_main_agent_forward <= forward_context, 'Provide valid minimum context for main agent.' + if batch_per_agent: # fetch frame-by-frame for agent + backward_context = forward_context = 0 + SynchronizedSceneDataset.set_context(self, backward=backward_context, forward=forward_context) + self.min_main_agent_forward = min_main_agent_forward if min_main_agent_forward else forward_context + self.min_main_agent_backward = min_main_agent_backward if min_main_agent_forward else backward_context + + # Extract all scenes from the scene dataset JSON for the appropriate split + scenes = BaseDataset._extract_scenes_from_scene_dataset_json( + dataset_json, split=split, requested_autolabels=requested_autolabels, is_datums_synchronized=True + ) + metadata = BaseDataset._extract_metadata_from_scene_dataset_json(dataset_json) + + # Return SynchronizedDataset with scenes built from dataset.json + dataset_metadata = DatasetMetadata.from_scene_containers(scenes, requested_annotations, requested_autolabels) + name_to_id = dataset_metadata.ontology_table[self.annotation_reference].name_to_id + self.requested_main_agent_types = tuple([name_to_id[atype] + 1 for atype in requested_main_agent_types]) + self.requested_main_agent_attributes = requested_main_agent_attributes + + # Resample scenes based on given fps + self.sampling_rate = fps / metadata.frame_per_second if fps != -1 and metadata.frame_per_second else 1.0 + assert self.sampling_rate <= 1, f"Support lower fps only (current is {metadata.frame_per_second:.2f} fps)." + resampled_scenes = [] + for scene in scenes: + resampled_scene = ResampledSceneContainer( + scene_path=scene.scene_path, + directory=scene.directory, + autolabeled_scenes=scene.autolabeled_scenes, + is_datums_synchronized=scene.is_datums_synchronized, + use_diskcache=scene.use_diskcache, + sampling_rate=self.sampling_rate + ) + resampled_scenes.append(resampled_scene) + + super().__init__( + dataset_metadata, + scenes=resampled_scenes, + datum_names=datum_names, + requested_annotations=requested_annotations + ) + + # Record each agent's life time + self.batch_per_agent = batch_per_agent + if batch_per_agent: + self.dataset_agent_index = defaultdict(list) + for index in range(len(self.dataset_item_index)): + scene_idx, sample_idx_in_scene, main_agent_info, datum_names = self.dataset_item_index[index] + _, main_agent_id = main_agent_info + # Find the range of agents' trajectories + if main_agent_id not in self.dataset_agent_index: + self.dataset_agent_index[main_agent_id] = [-1, -1, float('inf'), -1, []] + self.dataset_agent_index[main_agent_id] = [ + main_agent_id, + scene_idx, + min(sample_idx_in_scene, self.dataset_agent_index[main_agent_id][2]), # birth sample index + max(sample_idx_in_scene, self.dataset_agent_index[main_agent_id][3]), # death sample item index + datum_names + ] + self.dataset_agent_index = [v for k, v in self.dataset_agent_index.items()] + + def _build_item_index(self): + """Builds an index of dataset items that refer to the scene index, agent index, + sample index and datum_within_scene index. This refers to a particular dataset + split. __getitem__ indexes into this look up table. + + Synchronizes at the sample-level and only adds sample indices if context frames are available. + This is enforced by adding sample indices that fall in (bacwkard_context, N-forward_context) range. + + Returns + ------- + item_index: list + List of dataset items that contain index into + (scene_idx, sample_idx_in_scene, (main_agent_idx, main_agent_id), [datum_name ...]). + """ + logging.info(f'Building index for {self.__class__.__name__}, this will take a while.') + st = time.time() + # Fetch the item index per scene based on the selected datums. + with Pool(cpu_count()) as proc: + item_index = proc.starmap( + self._item_index_for_scene, [(scene_idx, ) for scene_idx in range(len(self.scenes))] + ) + logging.info(f'Index built in {time.time() - st:.2f}s.') + assert len(item_index) > 0, 'Failed to index items in the dataset.' + # Chain the index across all scenes. + item_index = list(itertools.chain.from_iterable(item_index)) + # Filter out indices that failed to load. + item_index = [item for item in item_index if item is not None] + item_lengths = [len(item_tup) for item_tup in item_index] + assert all([l == item_lengths[0] for l in item_lengths] + ), ('All sample items are not of the same length, datum names might be missing.') + return item_index + + def _item_index_for_scene(self, scene_idx): + scene = self.scenes[scene_idx] + instance_id_to_trajectory = defaultdict(list) + instance_id_to_segment_idx = defaultdict(int) + # Decide main trajectories reference. + # There are 3 cases to decide referenced annotation for trajectories: + # 1. LIDAR only: bounding_box_3d as reference. + # 2. CAMERA only: bounding_box_3d if use_3d_trajectories else bounding_box_2d. + # 3. LIDAR + CAMERA: bounding_box_3d (from LIDAR) as reference. + reference_datums = [datum_name for datum_name in scene.selected_datums \ + if self.trajectories_reference in datum_name] + + # Only add to index if datum-name exists. + if len(reference_datums) == 0: + logging.debug('Skipping scene {} due to missing datums'.format(scene)) + return [] + + # Define a safe sample range given desired context + sample_range = np.arange(0, len(scene.datum_index)) + annotated_samples = scene.annotation_index[scene.datum_index[sample_range]].any(axis=(1, 2)) + for sample_idx_in_scene, is_annotated in zip(sample_range, annotated_samples): + if not is_annotated: + continue + else: + for datum_name in reference_datums: + datum = self.get_datum(scene_idx, sample_idx_in_scene, datum_name) + annotation_key = self.annotation_reference + annotations = self.get_annotations(datum) + annotation_file = os.path.join( + self.scenes[scene_idx].directory, annotations[ANNOTATION_KEY_TO_TYPE_ID[annotation_key]] + ) + bboxes_list = ANNOTATION_REGISTRY[annotation_key].load( + annotation_file, self.dataset_metadata.ontology_table[annotation_key] + ) + for agent_idx, bbox in enumerate(bboxes_list): + # Filter undesired agent types and attributes. + if bbox.class_id not in self.requested_main_agent_types or \ + (self.ATTRIBUTE_NAME in bbox.attributes and \ + self.requested_main_agent_attributes is not None and \ + bbox.attributes[self.ATTRIBUTE_NAME] not in self.requested_main_agent_attributes): + continue + # Make sure the track is sample-continuous + instance_index_prefix = \ + f'{datum.id.name}_{str(scene_idx)}_{str(bbox.instance_id)}' + segment_idx_start = instance_id_to_segment_idx[instance_index_prefix] \ + if instance_index_prefix in instance_id_to_segment_idx else 0 + for segment_idx in range(segment_idx_start, len(self.scenes[scene_idx].samples)): + instance_index_id = f'{instance_index_prefix}_{segment_idx}' + if instance_index_id in instance_id_to_trajectory and \ + instance_id_to_trajectory[instance_index_id][-1][1] + 1 != sample_idx_in_scene: + continue + instance_id_to_trajectory[instance_index_id].append( + (scene_idx, sample_idx_in_scene, (agent_idx, bbox.instance_id), scene.selected_datums) + ) + instance_id_to_segment_idx[instance_index_prefix] = segment_idx + break + # Fiter unavailable items according to forward_context/backward_context for each agent. + item_index = [] + trajectory_min_length = self.min_main_agent_backward + self.min_main_agent_forward + 1 + for id_ in instance_id_to_trajectory: + scene_idx = instance_id_to_trajectory[id_][0][0] + num_samples = len(self.scenes[scene_idx].samples) + trajectory_length = len(instance_id_to_trajectory[id_]) + # Make sure track length is sufficient + if trajectory_length >= trajectory_min_length: + first_sample_idx = instance_id_to_trajectory[id_][0][1] + final_sample_idx = instance_id_to_trajectory[id_][-1][1] + # Crop out valid samples as items + beg = self.min_main_agent_backward \ + if self.min_main_agent_backward + first_sample_idx > self.backward_context \ + else self.backward_context + end = trajectory_length - (self.min_main_agent_forward \ + if self.min_main_agent_forward + final_sample_idx < num_samples \ + else self.forward_context) + if end > beg: + item_index.append(instance_id_to_trajectory[id_][beg:end]) + return list(itertools.chain.from_iterable(item_index)) + + def __len__(self): + """Return the length of the dataset.""" + return len(self.dataset_agent_index) if self.batch_per_agent else len(self.dataset_item_index) + + def __getitem__(self, index): + """Get the dataset item at index. + + Parameters + ---------- + index: int + Index of item to get. + + Returns + ------- + data: list of list of OrderedDict + + "timestamp": int + Timestamp of the image in microseconds. + + "datum_name": str + Sensor name from which the data was collected + + "rgb": PIL.Image (mode=RGB) + Image in RGB format. + + "intrinsics": np.ndarray + Camera intrinsics. + + "extrinsics": Pose + Camera extrinsics with respect to the world frame. + + "pose": Pose + Pose of sensor with respect to the world/global frame + + "main_agent_idx": int + Index of main agent in agent list (bounding_box_2d/bounding_box_3d). + + Returns a list of list of OrderedDict(s). + Outer list corresponds to temporal ordering of samples. Each element is + a list of OrderedDict(s) corresponding to synchronized datums. + + In other words, __getitem__ returns a nested list with the ordering as + follows: (C, D, I), where + C = forward_context + backward_context + 1, + D = len(datum_names) + I = OrderedDict item + """ + assert self.dataset_item_index is not None, ('Index is not built, select datums before getting elements.') + + if self.batch_per_agent: + # Get dataset agent index + main_agent_id, scene_idx, sample_idx_in_scene_start, sample_idx_in_scene_end, datum_names = \ + self.dataset_agent_index[index] + else: + # Get dataset item index + scene_idx, sample_idx_in_scene, main_agent_info, datum_names = self.dataset_item_index[index] + _, main_agent_id = main_agent_info + sample_idx_in_scene_start = sample_idx_in_scene - self.backward_context + sample_idx_in_scene_end = sample_idx_in_scene + self.forward_context + + # All sensor data (including pose, point clouds and 3D annotations are + # defined with respect to the sensor's reference frame captured at that + # corresponding timestamp. In order to move to a locally consistent + # reference frame, you will need to use the "pose" that specifies the + # ego-pose of the sensor with respect to the local (L) frame (pose_LS). + + context_window = [] + reference_annotation_key = self.annotation_reference + # Iterate through context samples + for qsample_idx_in_scene in range(sample_idx_in_scene_start, sample_idx_in_scene_end + 1): + # Main agent index may be different along the samples. + synchronized_sample = [] + for datum_name in datum_names: + datum_data = SynchronizedSceneDataset.get_datum_data(self, scene_idx, qsample_idx_in_scene, datum_name) + + if reference_annotation_key in datum_data: + # Over the main agent's trajectory, set main_agent_idx as None. + # Notice: main agent index may be different along the samples. + instance_matched = [ + bbox.instance_id == main_agent_id for bbox in datum_data[reference_annotation_key] + ] + main_agent_idx_in_sample = instance_matched.index(True) if any(instance_matched) else None + datum_data['main_agent_idx'] = main_agent_idx_in_sample + + synchronized_sample.append(datum_data) + context_window.append(synchronized_sample) + return context_window diff --git a/dgp/scripts/backfill_agents.py b/dgp/scripts/backfill_agents.py new file mode 100644 index 00000000..3df1f40f --- /dev/null +++ b/dgp/scripts/backfill_agents.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python +# Copyright 2021 Toyota Research Institute. All rights reserved. +"""Script that backfill agents into original DGP format.""" +import argparse +import logging +import os +from collections import defaultdict +from pathlib import Path + +from dgp import ( + AGENT_FOLDER, FEATURE_ONTOLOGY_FOLDER, TRI_DGP_AGENT_TRACKS_JSON_NAME, TRI_DGP_AGENTS_JSON_NAME, + TRI_DGP_AGENTS_SLICES_JSON_NAME +) +from dgp.datasets.prediction_dataset import PredictionAgentDataset +from dgp.proto import dataset_pb2, features_pb2 +from dgp.proto.agent_pb2 import (AgentGroup, AgentsSlice, AgentsSlices, AgentTracks) +from dgp.proto.ontology_pb2 import FeatureOntology, FeatureOntologyItem +from dgp.utils.dataset_conversion import get_date, get_datetime_proto +from dgp.utils.protobuf import (generate_uid_from_pbobject, open_pbobject, save_pbobject_as_json) + +TRIPCCOntology = FeatureOntology(items=[FeatureOntologyItem(name='ParkedVehicleState', id=0, feature_value_type=0)]) + + +class AgentBackfillDemo: + """ + Class to backfill agents information into ioda scene dataset to create a + demo dataset. + """ + def __init__(self, scene_dataset_json, creator): + self.agent_dataset_name = "agents_pcc_mini" + self.version = "1" + self.description = 'agents of pcc mini dataset' + self.EMAIL = creator + self.public = True + self.scene_dataset_json = scene_dataset_json + self.scene_dataset = open_pbobject(scene_dataset_json, dataset_pb2.SceneDataset) + self.agents_dataset_pb2 = dataset_pb2.Agents() + self.local_output_path = Path(scene_dataset_json).parent.as_posix() + self.ontologies = {features_pb2.PARKED_CAR: TRIPCCOntology} + self.populate_metadata() + + def populate_metadata(self): + """Populate boilerplate fields agent metadata""" + self.agents_dataset_pb2.metadata.name = self.agent_dataset_name + self.agents_dataset_pb2.metadata.version = self.version + self.agents_dataset_pb2.metadata.creation_date = get_date() + self.agents_dataset_pb2.metadata.creator = self.EMAIL + self.agents_dataset_pb2.metadata.description = self.description + self.agents_dataset_pb2.metadata.origin = self.agents_dataset_pb2.metadata.PUBLIC if self.public \ + else self.agents_dataset_pb2.metadata.INTERNAL + + def generate(self): + for (split_type, split) in zip([dataset_pb2.TRAIN, dataset_pb2.TEST, dataset_pb2.VAL], + ['train', 'test', 'val']): + if split_type not in self.scene_dataset.scene_splits: + continue + + logging.info('Processing {}'.format(split)) + original_dataset = PredictionAgentDataset( + self.scene_dataset_json, + split=split, + datum_names=('LIDAR', ), + requested_main_agent_types=( + 'Person', + 'Truck', + 'Car', + 'Bus/RV/Caravan', + 'Motorcycle', + 'Wheeled Slow', + 'Train', + 'Towed Object', + 'Bicycle', + 'Trailer', + ), + batch_per_agent=True + ) + self.backfill(original_dataset, split_type) + agent_json_path = self.write_agents() + + return agent_json_path + + def backfill(self, original_dataset, split_type): + """Backfill agent information. + + Parameters + ---------- + original_dataset: PredictionAgentDataset + DGP PredictionAgentDataset object + + split_type: DatasetSplit + Split of dataset to read ("train" | "val" | "test" | "train_overfit"). + + """ + # Map from scene idx to list agent + scene_agent_map = defaultdict(list) + for agent_idx, agent_track in enumerate(original_dataset.dataset_agent_index): + scene_idx = agent_track[1] + scene_agent_map[scene_idx].append(agent_idx) + for scene_idx, scene in enumerate(original_dataset.scenes): + output_scene_dirname = scene.directory + + agent_pb2 = AgentGroup() + agent_pb2.feature_ontologies[features_pb2.PARKED_CAR] = \ + generate_uid_from_pbobject(self.ontologies[features_pb2.PARKED_CAR]) + for k, v in scene.scene.ontologies.items(): + agent_pb2.agent_ontologies[k] = v + agent_tracks_pb2 = AgentTracks() + agents_slices_pb2 = AgentsSlices() + sample_agent_snapshots_dict = defaultdict(AgentsSlice) + for agent_idx in scene_agent_map[scene_idx]: + main_agent_id, scene_idx, sample_idx_in_scene_start, _, _ = \ + original_dataset.dataset_agent_index[agent_idx] + agent_track_original = original_dataset[agent_idx] + agent_track = agent_tracks_pb2.agent_tracks.add() + agent_track.instance_id = main_agent_id + agent_track.class_id = original_dataset.dataset_metadata.ontology_table[ + original_dataset.annotation_reference].contiguous_id_to_class_id[agent_track_original[0][0][ + 'bounding_box_3d'][agent_track_original[0][0]['main_agent_idx']].class_id] + sample_idx = sample_idx_in_scene_start + for agent_snapshot_original in agent_track_original: + try: + box = agent_snapshot_original[0]['bounding_box_3d'][int( + agent_snapshot_original[0]['main_agent_idx'] + )] + except TypeError: # pylint: disable=bare-except + sample_idx = sample_idx + 1 + continue + if sample_idx not in sample_agent_snapshots_dict: + sample_agent_snapshots_dict[sample_idx].slice_id.CopyFrom(scene.samples[sample_idx].id) + sample_agent_snapshots_dict[sample_idx].slice_id.index = sample_idx + agent_snapshot = agent_track.agent_snapshots.add() + agent_snapshot.agent_snapshot_3D.box.CopyFrom(box.to_proto()) + agent_snapshot.slice_id.CopyFrom(scene.samples[sample_idx].id) + agent_snapshot.slice_id.index = sample_idx + agent_snapshot.agent_snapshot_3D.class_id = agent_track.class_id + + agent_snapshot.agent_snapshot_3D.instance_id = main_agent_id + # The feature fields need to backfill from + agent_snapshot.agent_snapshot_3D.features.extend(["dynamic"]) + + # 5 for parked car features + agent_snapshot.agent_snapshot_3D.feature_type = features_pb2.PARKED_CAR + + # Handle agent slices + sample_agent_snapshots_dict[sample_idx].agent_snapshots.extend([agent_snapshot]) + + sample_idx = sample_idx + 1 + + for sample_idx in range(len(scene.samples)): + if sample_idx in sample_agent_snapshots_dict: + agents_slices_pb2.agents_slices.extend([sample_agent_snapshots_dict[sample_idx]]) + else: + agents_slices_pb2.agents_slices.extend([AgentsSlice()]) + + # Save agent_tracks file + os.makedirs(os.path.join(output_scene_dirname, AGENT_FOLDER), exist_ok=True) + agent_tracks_filename = os.path.join( + AGENT_FOLDER, + TRI_DGP_AGENT_TRACKS_JSON_NAME.format(track_hash=generate_uid_from_pbobject(agents_slices_pb2)) + ) + + save_pbobject_as_json(agent_tracks_pb2, os.path.join(output_scene_dirname, agent_tracks_filename)) + agent_pb2.agent_tracks_file = agent_tracks_filename + + # Save agents_slices file + agents_slices_filename = os.path.join( + AGENT_FOLDER, + TRI_DGP_AGENTS_SLICES_JSON_NAME.format(slice_hash=generate_uid_from_pbobject(agent_tracks_pb2)) + ) + save_pbobject_as_json(agents_slices_pb2, os.path.join(output_scene_dirname, agents_slices_filename)) + agent_pb2.agents_slices_file = agents_slices_filename + + # Save AgentGroup + agent_pb2.log = scene.scene.log + agent_pb2.name = scene.scene.name + agent_pb2.creation_date.CopyFrom(get_datetime_proto()) + + scene_dirname = output_scene_dirname.replace(self.local_output_path + '/', '') + agents_group_filename = os.path.join( + scene_dirname, TRI_DGP_AGENTS_JSON_NAME.format(agent_hash=generate_uid_from_pbobject(agent_pb2)) + ) + + self.agents_dataset_pb2.agents_splits[split_type].filenames.append(agents_group_filename) + save_pbobject_as_json(agent_pb2, os.path.join(self.local_output_path + '/' + agents_group_filename)) + + # Populate ontologies + os.makedirs(os.path.join(output_scene_dirname, FEATURE_ONTOLOGY_FOLDER), exist_ok=True) + for feature_type, ontology_id in agent_pb2.feature_ontologies.items(): + ontology_filename = os.path.join( + output_scene_dirname, FEATURE_ONTOLOGY_FOLDER, "{}.json".format(ontology_id) + ) + save_pbobject_as_json(self.ontologies[feature_type], ontology_filename) + + def write_agents(self, ): + """Write the final scene dataset JSON. + Returns + ------- + scene_dataset_json_path: str + Path of the scene dataset JSON file created. + """ + agent_dataset_json_path = os.path.join( + self.local_output_path, + '{}_v{}.json'.format(self.agents_dataset_pb2.metadata.name, self.agents_dataset_pb2.metadata.version) + ) + save_pbobject_as_json(self.agents_dataset_pb2, agent_dataset_json_path) + + # Printing SceneDataset scene counts per split (post-merging) + logging.info('-' * 80) + logging.info( + 'Output SceneDataset {} has: {} train, {} val, {} test'.format( + agent_dataset_json_path, len(self.agents_dataset_pb2.agents_splits[dataset_pb2.TRAIN].filenames), + len(self.agents_dataset_pb2.agents_splits[dataset_pb2.VAL].filenames), + len(self.agents_dataset_pb2.agents_splits[dataset_pb2.TEST].filenames) + ) + ) + + return agent_dataset_json_path + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, add_help=True + ) + parser.add_argument("-i", "--scene-dataset-json", help="Input dataset json", required=True) + parser.add_argument('--verbose', action='store_true') + parser.add_argument("--creator", help="Creator email", required=True) # + args, other_args = parser.parse_known_args() + if args.verbose: + logging.basicConfig(level=logging.INFO) + + dataset = AgentBackfillDemo(args.scene_dataset_json, args.creator) + dataset.generate() diff --git a/dgp/utils/visualization_utils.py b/dgp/utils/visualization_utils.py index c43c23df..61dc6e35 100644 --- a/dgp/utils/visualization_utils.py +++ b/dgp/utils/visualization_utils.py @@ -385,6 +385,12 @@ class BEVImage: background_clr: tuple[int], defaults: (0, 0, 0) Background color in BGR order. + + center_offset_w: int, default: 0 + Offset in pixels to move ego center in BEV. + + center_offset_h: int, default: 0 + Offset in pixels to move ego center in BEV. """ def __init__( self, @@ -394,7 +400,9 @@ def __init__( polar_step_size_meters=10, forward=(1, 0, 0), left=(0, 1, 0), - background_clr=(0, 0, 0) + background_clr=(0, 0, 0), + center_offset_w=0, + center_offset_h=0, ): forward, left = np.array(forward, np.float64), np.array(left, np.float64) assert np.dot(forward, left) == 0 # orthogonality check. @@ -405,7 +413,7 @@ def __init__( self._polar_step_size_meters = polar_step_size_meters self._forward = forward self._left = left - self._bg_clr = np.uint8(background_clr)[::-1].reshape(1, 1, 3) + self._bg_clr = np.array(background_clr)[::-1].reshape(1, 1, 3).astype(np.uint8) # Body frame -> BEV frame right = -left @@ -413,8 +421,10 @@ def __init__( bev_rotation = Pose.from_rotation_translation(bev_rotation, tvec=np.zeros(3)) self._bev_rotation = bev_rotation - self._center_pixel = (int(metric_height * pixels_per_meter) // 2, int(metric_width * pixels_per_meter) // 2) - + self._center_pixel = ( + int((metric_width * pixels_per_meter) // 2 - pixels_per_meter * center_offset_w), + int((metric_height * pixels_per_meter) // 2 - pixels_per_meter * center_offset_h) + ) self.reset() def __repr__(self): @@ -432,7 +442,7 @@ def reset(self): for i in range(1, int(max(self._metric_width, self._metric_height)) // self._polar_step_size_meters): cv2.circle( self.data, self._center_pixel, int(i * self._polar_step_size_meters * self._pixels_per_meter), - (50, 50, 50), 2 + (50, 50, 50), 1 ) def render_point_cloud(self, point_cloud, extrinsics=Pose(), color=GRAY): @@ -544,6 +554,46 @@ def clip_norm(v, x): color = (255, 110, 199) cv2.arrowedLine(self.data, (cx, cy), (cx2, cy2), color, thickness=1, line_type=cv2.LINE_AA) + def render_paths(self, paths, extrinsics=Pose(), colors=(GREEN, ), line_thickness=1, tint=1.0): + """Render object paths on bev. + + Parameters + ---------- + paths: list[list[Pose]] + List of object poses in the coordinate frame of the current timestep. + + extrinsics: Pose, default: Identity pose + The pose of the pointcloud sensor wrt the body frame (Sensor frame -> (Vehicle) Body frame). + + colors: List of RGB tuple, default: [GREEN,] + Draw path using this color. + + line_thickness: int, default: 1 + Thickness of lines. + + tint: float, default: 1.0 + Mulitiplicative factor applied to color used to darken lines. + """ + + if len(colors) == 1: + colors = list(colors) * len(paths) + + if tint != 1.0: + colors = [[int(tint * c) for c in color] for color in colors] + + combined_transform = self._bev_rotation * extrinsics + + for path, color in zip(paths, colors): + # path should contain a list of Pose objects or None types. None types will be skipped. + # TODO: add option to interpolate skipped poses. + path3d = [combined_transform * pose.tvec.reshape(1, 3) for pose in path if pose is not None] + path2d = np.round(self._pixels_per_meter * np.stack(path3d, 0)[..., :2], + 0).astype(np.int32).reshape(1, -1, 2) + offset = np.array(self._center_pixel).reshape(1, 1, 2) # pylint: disable=E1121 + path2d = path2d + offset + # TODO: if we group the paths by color we can draw all paths with the same color at once + cv2.polylines(self.data, path2d, 0, color, line_thickness, cv2.LINE_AA) + def render_bounding_box_3d( self, bboxes3d, @@ -554,12 +604,16 @@ def render_bounding_box_3d( texts=None, line_thickness=2, font_scale=0.5, + font_colors=(WHITE, ), + markers=None, + marker_scale=.5, + marker_colors=(RED, ), ): """Render bounding box 3d in BEV perspective. Parameters ---------- - bboxes3d: list of BoundingBox3D + bboxes3d: List of BoundingBox3D 3D annotations in the sensor coordinate frame. extrinsics: Pose, default: Identity pose @@ -578,21 +632,40 @@ def render_bounding_box_3d( 3D annotation category name. line_thickness: int, default: 2 - Thickness of lines + Thickness of lines. font_scale: float, default: 0.5 Font scale used for text labels. + + font_colors: List of RGB tuple, default: [WHITE,] + Color used for text labels. + + markers: List[int], default: None + List of opencv markers to draw in bottom right corner of cuboid. Should be one of: + cv2.MARKER_CROSS, cv2.MARKER_DIAMOND, cv2.MARKER_SQUARE, cv2.MARKER_STAR, cv2.MARKER_TILTED_CROSS, cv2.MARKER_TRIANGLE_DOWN, cv2.MARKER_TRIANGLE_UP, or None. + + marker_scale: float, default: .5 + Scale factor for markers, + + marker_colors: List of RGB Tuple, default: [RED,] + Draw markers using this color. """ if len(colors) == 1: colors = list(colors) * len(bboxes3d) + if len(font_colors) == 1: + font_colors = list(font_colors) * len(bboxes3d) + + if len(marker_colors) == 1: + marker_colors = list(marker_colors) * len(bboxes3d) + combined_transform = self._bev_rotation * extrinsics # Draw cuboids for bidx, (bbox, color) in enumerate(zip(bboxes3d, colors)): # Create 3 versions of colors for face coding. - front_face_color = RED + front_face_color = color side_line_color = [int(side_color_fraction * c) for c in color] rear_face_color = [int(rear_color_fraction * c) for c in color] @@ -608,26 +681,41 @@ def render_bounding_box_3d( center = np.mean(corners2d, axis=0).astype(np.int32) corners2d = corners2d.astype(np.int32) + # Draw front face, side faces and back face + cv2.line(self.data, tuple(corners2d[0]), tuple(corners2d[1]), front_face_color, line_thickness, cv2.LINE_AA) + cv2.line(self.data, tuple(corners2d[1]), tuple(corners2d[2]), side_line_color, line_thickness, cv2.LINE_AA) + cv2.line(self.data, tuple(corners2d[2]), tuple(corners2d[3]), rear_face_color, line_thickness, cv2.LINE_AA) + cv2.line(self.data, tuple(corners2d[3]), tuple(corners2d[0]), side_line_color, line_thickness, cv2.LINE_AA) + # Draw white light connecting center and font side. - cv2.line( + cv2.arrowedLine( self.data, tuple(center), ( (corners2d[0][0] + corners2d[1][0]) // 2, (corners2d[0][1] + corners2d[1][1]) // 2, - ), GREEN, 2 + ), WHITE, 1, cv2.LINE_AA ) - # Draw front face, side faces and back face - cv2.line(self.data, tuple(corners2d[0]), tuple(corners2d[1]), front_face_color, line_thickness) - cv2.line(self.data, tuple(corners2d[1]), tuple(corners2d[2]), side_line_color, line_thickness) - cv2.line(self.data, tuple(corners2d[2]), tuple(corners2d[3]), rear_face_color, line_thickness) - cv2.line(self.data, tuple(corners2d[3]), tuple(corners2d[0]), side_line_color, line_thickness) - if texts: - top_left = np.argmin(np.linalg.norm(corners2d, axis=1)) - cv2.putText( - self.data, texts[bidx], tuple(corners2d[top_left]), cv2.FONT_HERSHEY_SIMPLEX, font_scale, WHITE, - line_thickness // 2, cv2.LINE_AA - ) + if texts[bidx] is not None: + top_left = np.argmin(np.linalg.norm(corners2d, axis=1)) + cv2.putText( + self.data, texts[bidx], tuple(corners2d[top_left]), cv2.FONT_HERSHEY_SIMPLEX, font_scale, + font_colors[bidx], line_thickness // 2, cv2.LINE_AA + ) + + if markers: + if markers[bidx] is not None: + bottom_right = np.argmax(np.linalg.norm(corners2d, axis=1)) + + assert markers[bidx] in [ + cv2.MARKER_CROSS, cv2.MARKER_DIAMOND, cv2.MARKER_SQUARE, cv2.MARKER_STAR, + cv2.MARKER_TILTED_CROSS, cv2.MARKER_TRIANGLE_DOWN, cv2.MARKER_TRIANGLE_UP + ] + + cv2.drawMarker( + self.data, tuple(corners2d[bottom_right]), marker_colors[bidx], markers[bidx], + int(20 * marker_scale), 2, cv2.LINE_AA + ) def render_camera_frustrum(self, intrinsics, extrinsics, width, color=YELLOW, line_thickness=1): """ @@ -774,6 +862,13 @@ def visualize_bev( bev_line_thickness=4, bev_font_scale=0.5, radar_datums=None, + instance_colormap=None, + cuboid_caption_fn=None, + marker_fn=None, + marker_scale=.5, + show_paths_on_bev=False, + bev_center_offset_w=0, + bev_center_offset_h=0, ): """Create BEV visualization that shows pointcloud, 3D bounding boxes, and optionally camera frustrums. Parameters @@ -782,7 +877,7 @@ def visualize_bev( List of lidar datums as a dictionary. class_colormap: Dict Mapping from class IDs to RGB colors. - show_instance_id_on_bev: bool, default: False + show_instance_id_on_bev: Bool, default: False If True, then show `instance_id` on a corner of 3D bounding boxes in BEV view. If False, then show `class_name` instead. id_to_name: OrderedDict, default: None @@ -795,14 +890,39 @@ def visualize_bev( See `BEVImage` for these keyword arguments. radar_datums: List[OrderedDict], default: None List of radar datums to visualize + instance_colormap: Dict + Mapping from instance id to RGB colors. + cuboid_caption_fn: Callable, BoundingBox3d -> Tuple[String,Tuple[3]] + Function taking a BoundingBox3d object and returning a tuple with the caption string, and the rgb + value for that caption. e.g., ( 'car', (255,0,0) ) + marker_fn: Callable, BoundingBox3d -> Tuple[int,Tuple[3]] + Function taking a BoundingBox3d object and returning a tuple with the caption a marker id, and the rgb + value for that marker. e.g., ( cv2.MARKER_DIAMOND, (255,0,0) ). Marker should be one of + cv2.MARKER_CROSS, cv2.MARKER_DIAMOND, cv2.MARKER_SQUARE, cv2.MARKER_STAR, cv2.MARKER_TILTED_CROSS, cv2.MARKER_TRIANGLE_DOWN, cv2.MARKER_TRIANGLE_UP, or None. + show_paths_on_bev: Bool, default: False + If true draw a path for each cuboid. Paths are stored in cuboid attributes under the 'path' key, i.e., + path = cuboid.attributes['path'], paths are themselves a list of pose objects transformed to the + correct frame. This method does not handle creating or transforming the paths. + bev_enter_offset_w: int, default: 0 + Offset in pixels to move ego center in BEV. + bev_center_offset_h: int, default: 0 + Offset in pixels to move ego center in BEV. + Returns ------- np.ndarray BEV visualization as an image. """ bev = BEVImage( - bev_metric_width, bev_metric_height, bev_pixels_per_meter, bev_polar_step_size_meters, bev_forward, bev_left, - bev_background_clr + bev_metric_width, + bev_metric_height, + bev_pixels_per_meter, + bev_polar_step_size_meters, + bev_forward, + bev_left, + bev_background_clr, + center_offset_w=bev_center_offset_w, + center_offset_h=bev_center_offset_h ) # 1. Render pointcloud @@ -823,24 +943,58 @@ def visualize_bev( # 3. Render 3D bboxes. for lidar_datum in lidar_datums: if 'bounding_box_3d' in lidar_datum: - class_ids = [bbox3d.class_id for bbox3d in lidar_datum['bounding_box_3d']] - colors = [class_colormap[class_id] for class_id in class_ids] - if show_instance_id_on_bev: - labels = [str(bbox3d.instance_id) for bbox3d in lidar_datum['bounding_box_3d']] - else: # show class names - labels = [id_to_name[i] for i in class_ids] + + if len(lidar_datum['bounding_box_3d']) == 0: + continue + + if instance_colormap is not None: + colors = [ + instance_colormap.get(bbox.instance_id, class_colormap[bbox.class_id]) + for bbox in lidar_datum['bounding_box_3d'] + ] + else: + colors = [class_colormap[bbox.class_id] for bbox in lidar_datum['bounding_box_3d']] + + # If no caption function is supplied, generate one from the instance ids or class ids + # Caption functions should return a tuple (string, color) + # TODO: expand to include per caption font size. + if show_instance_id_on_bev and cuboid_caption_fn is None: + cuboid_caption_fn = lambda x: (str(x.instance_id), WHITE) + elif cuboid_caption_fn is None: # show class names + cuboid_caption_fn = lambda x: (id_to_name[x.class_id], WHITE) + + labels, font_colors = zip(*[cuboid_caption_fn(bbox3d) for bbox3d in lidar_datum['bounding_box_3d']]) + + markers, marker_colors = None, (RED, ) + if marker_fn is not None: + markers, marker_colors = zip(*[marker_fn(bbox3d) for bbox3d in lidar_datum['bounding_box_3d']]) + bev.render_bounding_box_3d( lidar_datum['bounding_box_3d'], lidar_datum['extrinsics'], colors=colors, - texts=labels, + texts=labels if bev_font_scale > 0 else None, line_thickness=bev_line_thickness, font_scale=bev_font_scale, + font_colors=font_colors, + markers=markers if marker_scale > 0 else None, + marker_scale=marker_scale, + marker_colors=marker_colors, ) + if show_paths_on_bev: + # Collect the paths and path colors + paths, path_colors = zip( + *[(bbox.attributes['path'], c) + for bbox, c in zip(lidar_datum['bounding_box_3d'], colors) + if 'path' in bbox.attributes] + ) + if len(paths) > 0: + bev.render_paths(paths, extrinsics=lidar_datum['extrinsics'], colors=path_colors, line_thickness=1) + + # 4. Render camera frustrums. if camera_datums is not None: for cam_datum, cam_color in zip(camera_datums, camera_colors): - # 3. Render camera frustrums. bev.render_camera_frustrum( cam_datum['intrinsics'], cam_datum['extrinsics'], diff --git a/examples/agent_visualization.py b/examples/agent_visualization.py new file mode 100644 index 00000000..1e8ab5ba --- /dev/null +++ b/examples/agent_visualization.py @@ -0,0 +1,168 @@ +# Copyright 2021 Toyota Research Institute. All rights reserved. +import argparse + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import seaborn as sns +from moviepy.editor import ImageSequenceClip +from tqdm import tqdm + +from dgp.constants import (vehicle_applanix_origin_to_r_bumper, vehicle_height, vehicle_length, vehicle_width) +from dgp.datasets.agent_dataset import AgentDatasetLite +from dgp.utils.pose import Pose +from dgp.utils.structures.bounding_box_3d import BoundingBox3D +from dgp.utils.visualization_utils import visualize_bev + + +def calc_warp_pose(pose_other, pose_target): + # return pose for going from pose to pose_target + pose_target_w = pose_target.inverse() # world to target + pose_p2p1 = pose_target_w * pose_other # local other to world, world to target -> local to target + return pose_p2p1 + + +def render_agent_bev( + agent_dataset, + ontology, +): + class_colormap = ontology._contiguous_id_colormap + id_to_name = ontology.contiguous_id_to_name + + tvec = np.array([vehicle_length / 2 - vehicle_applanix_origin_to_r_bumper, 0, 0]) + ego_box = BoundingBox3D( + Pose(tvec=tvec), sizes=np.array([vehicle_width, vehicle_length, vehicle_height]), class_id=1, instance_id=0 + ) + + # Drawing code, create a pallet + pallet = list(sns.color_palette("hls", 32)) + pallet = [[np.int(255 * a), np.int(255 * b), np.int(255 * c)] for a, b, c in pallet] + + def get_random_color(): + idx = np.random.choice(len(pallet)) + return pallet[idx] + + trackid_to_color = {} + paths = {} + frames = [] + prior_pose = None + max_path_len = 15 + + for k in tqdm(range(0, len(agent_dataset))): + context = agent_dataset[k] + lidar = context[0]["datums"][1] + camera = context[0]["datums"][0] + agents = context[0]['agents'] + cam_color = (0, 255, 0) + agents.boxlist.append(ego_box) + lidar['bounding_box_3d'] = agents + trackid_to_color[0] = (255, 255, 255) + + # core tracking color and path generation code + if prior_pose is None: + prior_pose = lidar['pose'] + + warp_pose = calc_warp_pose(prior_pose, lidar['pose']) + prior_pose = lidar['pose'] + + new_colors = {box.instance_id: get_random_color() for box in agents if box.instance_id not in trackid_to_color} + trackid_to_color.update(new_colors) + updated = [] + # warp existing paths + for instance_id in paths: + # move the path into ego's local frame. We assume all prior path entrys are in the previous frame + # this is not true if we skip a step because of occulision or missed detection... TODO: handle this somehow + paths[instance_id] = [warp_pose * pose if pose is not None else None for pose in paths[instance_id]] + + # add new boxes to the path + for box in agents: + + if box.instance_id not in paths: + paths[box.instance_id] = [] + + paths[box.instance_id].insert(0, box.pose) + + # keep track of what was updated so we can insert Nones if there is a miss + updated.append(box.instance_id) + + if len(paths[box.instance_id]) > max_path_len: + paths[box.instance_id].pop() + + box.attributes['path'] = paths[box.instance_id] + + # go through the non updated paths and append None + for instance_id in paths: + if instance_id not in updated: + paths[instance_id].insert(0, None) + + if len(paths[instance_id]) > max_path_len: + paths[instance_id].pop() + + cuboid_caption_fn = lambda x: ('Parked' if 'parked' in x.attributes else None, (255, 0, 0)) + + marker_fn = lambda x: (cv2.MARKER_DIAMOND if 'parked' in x.attributes else None, (255, 0, 0)) + + img = visualize_bev([lidar], class_colormap, show_instance_id_on_bev=False, id_to_name = id_to_name, bev_font_scale = .5, bev_line_thickness = 2\ + , instance_colormap = trackid_to_color, show_paths_on_bev=True,\ + cuboid_caption_fn = cuboid_caption_fn, \ + marker_fn = marker_fn, + camera_datums= [camera], + camera_colors = [cam_color], + bev_metric_width=100, + bev_metric_height=int(3*100/4), + bev_pixels_per_meter = 10, + bev_center_offset_w=0 + ) + + frames.append(img) + return frames + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, add_help=True + ) + + parser.add_argument('--agent-json', help='Agent JSON file') + parser.add_argument('--scene-dataset-json', help='Scene Dataset JSON file') + parser.add_argument( + '--split', help='Dataset split', required=False, choices=['train', 'val', 'test'], default='train' + ) + args, other_args = parser.parse_known_args() + + agent_dataset_frame = AgentDatasetLite( + args.scene_dataset_json, + args.agent_json, + split=args.split, + datum_names=['lidar', 'CAMERA_01'], + requested_main_agent_classes=('Car', 'Person'), + requested_feature_types=("parked_car", ), + batch_per_agent=False + ) + + agent_dataset_lite = AgentDatasetLite( + args.scene_dataset_json, + args.agent_json, + split=args.split, + datum_names=['lidar', 'CAMERA_01'], + requested_main_agent_classes=('Car', 'Person'), + requested_feature_types=("parked_car", ), + batch_per_agent=True + ) + + ont = agent_dataset_lite.Agent_dataset_metadata.ontology_table.get('bounding_box_3d', None) + + bev_frames = render_agent_bev(agent_dataset_frame, ont) + + a = [agent_dataset_frame.dataset_item_index[k][0] for k in range(len(agent_dataset_frame))] + + #Should design a better way + frame_num = 0 + for i in range(max(a) + 1): + + plt.figure(figsize=(20, 20)) + + clip = ImageSequenceClip(bev_frames[frame_num:frame_num + a.count(i)], fps=10) + clip.write_gif('test_scene' + str(i) + '.gif', fps=10) + frame_num += a.count(i) + print(frame_num)