Add ugv dataset and test for fov research

Author: spencer@primus

avapi/__init__.py

@@ -5,6 +5,7 @@
 import avapi.nuimages
 import avapi.nuscenes
 import avapi.opv2v
+import avapi.ugv
 import avapi.visualize
 
 

avapi/_dataset.py

@@ -131,6 +131,13 @@ def check_frame(self, frame):
             frame in self.frames
         ), f"Candidate frame, {frame}, not in frame set {self.frames}"
 
+    def get_agents(self, frame: int) -> "DataContainer":
+        return self._load_agents(frame)
+
+    def get_agent(self, frame: int, agent: int):
+        agents = self.get_agents(frame)
+        return [ag for ag in agents if ag.ID == agent][0]
+
     def get_agent_set(self, frame: int) -> set:
         return self._load_agent_set(frame=frame)
 
@@ -337,19 +344,25 @@ def save_objects(self, frame, objects, folder, file=None):
             os.makedirs(folder)
         self._save_objects(frame, objects, folder, file=file)
 
+    def _load_agents(self, frame):
+        raise NotImplementedError
+
+    def _load_ego(self, frame, agent=None):
+        raise NotImplementedError
+
     def _load_frames(self, sensor):
         raise NotImplementedError
 
     def _load_calibration(self, frame, sensor, reference):
         raise NotImplementedError
 
-    def _load_image(self, frame, camera, agent=None):
+    def _load_image(self, frame, sensor, agent=None):
         raise NotImplementedError
 
-    def _load_semseg_image(self, frame, camera, agent=None):
+    def _load_semseg_image(self, frame, sensor, agent=None):
         raise NotImplementedError
 
-    def _load_depth_image(self, frame, camera, agent=None):
+    def _load_depth_image(self, frame, sensor, agent=None):
         raise NotImplementedError
 
     def _load_lidar(self, frame, sensor, agent=None):
@@ -631,7 +644,7 @@ def __init__(
         self.nuX = nuX
         self.nuX_can = nuX_can
         if nuX_can is not None:
-            self.vehicle_pose = self.nuX_can.get_messages(self.scene["name"], "pose")
+            self.vehicle_pose = self.nuX_can.get_messages(self.scene, "pose")
             self.vehicle_pose_utime = np.array(
                 [vp["utime"] for vp in self.vehicle_pose]
             )
@@ -695,6 +708,10 @@ def _get_sensor_file_name(self, frame, sensor=None, agent=None):
     def _load_frames(self, **kwargs):
         return self.frames
 
+    def _load_agent_set(self, frame: int) -> set:
+        # TODO: this is slow...improve
+        return {ag.ID for ag in self.get_agents(frame)}
+
     def _load_calibration(self, frame, ego_reference, sensor=None, **kwargs):
         """
         W := global "world" frame
@@ -775,7 +792,7 @@ def _load_ego(self, frame, **kwargs):
         box3d = Box3D(x_G_to_E_in_G, q_G_to_E, self.hwl, ID=-1)
 
         # -- set up ego in global reference frame
-        veh = VehicleState(obj_type="car")
+        veh = VehicleState(obj_type="car", ID=1000)
         veh.set(
             t=t,
             position=x_G_to_E_in_G,

avapi/carla/dataset.py

@@ -285,13 +285,6 @@ def get_sensor_file(self, frame, timestamp, sensor, agent, file_type):
     def get_sensor_name(self, sensor, agent):
         return sensor
 
-    def get_agents(self, frame: int) -> "DataContainer":
-        return self._load_agents(frame)
-
-    def get_agent(self, frame: int, agent: int):
-        agents = self.get_agents(frame)
-        return [ag for ag in agents if ag.ID == agent][0]
-
     def _load_agent_set(self, frame: int) -> set:
         # TODO: this is slow...improve
         return {ag.ID for ag in self.get_agents(frame)}

avapi/nuscenes/dataset.py

@@ -4,6 +4,7 @@
 from typing import List
 
 import numpy as np
+from avstack.config import DATASETS
 from avstack.geometry.transformations import matrix_cartesian_to_spherical
 from scipy.interpolate import interp1d
 
@@ -26,10 +27,11 @@
     splits_scenes = None
 
 
+@DATASETS.register_module()
 class nuScenesManager(_nuManager):
     NAME = "nuScenes"
 
-    def __init__(self, data_dir, split="v1.0-mini", verbose=False):
+    def __init__(self, data_dir, split="v1.0-mini", max_scenes=None, verbose=False):
         nusc = NuScenes(version=split, dataroot=data_dir, verbose=verbose)
         try:
             nusc_can = NuScenesCanBus(dataroot=data_dir)
@@ -40,8 +42,11 @@ def __init__(self, data_dir, split="v1.0-mini", verbose=False):
         self.scene_name_to_index = {}
         self.scene_number_to_index = {}
         self.index_to_scene = {}
-        self.scenes = [sc["name"] for sc in nusc.scene]
-        for i, sc in enumerate(nusc.scene):
+        n_scenes = (
+            len(nusc.scene) if max_scenes is None else min(len(nusc.scene), max_scenes)
+        )
+        self.scenes = [sc["name"] for sc in nusc.scene][:n_scenes]
+        for i, sc in enumerate(nusc.scene[:n_scenes]):
             self.scene_name_to_index[sc["name"]] = i
             self.scene_number_to_index[int(sc["name"].replace("scene-", ""))] = i
             self.index_to_scene[i] = sc["name"]
@@ -128,13 +133,13 @@ def __init__(
         except Exception as e:
             logging.warning("Cannot find CAN bus data")
             nusc_can = None
-        self.scene = scene
-        self.scene_name = self.scene
+        self.scene = scene["name"]
+        self.sequence = scene  # a dictionary
         self.framerate = 2
         self.sequence_id = scene["name"]
         self.splits_scenes = splits_scenes
         try:
-            veh_speed = self.nuX_can.get_messages(self.scene["name"], "vehicle_monitor")
+            veh_speed = self.nuX_can.get_messages(self.scene, "vehicle_monitor")
         except Exception as e:
             self.ego_speed_interp = None
         else:
@@ -151,14 +156,17 @@ def __init__(
 
     def make_sample_records(self):
         self.sample_records = {
-            0: self.nuX.get("sample", self.scene["first_sample_token"])
+            0: self.nuX.get("sample", self.sequence["first_sample_token"])
         }
-        for i in range(1, self.scene["nbr_samples"], 1):
+        for i in range(1, self.sequence["nbr_samples"], 1):
             self.sample_records[i] = self.nuX.get(
                 "sample", self.sample_records[i - 1]["next"]
             )
         self.t0 = self.sample_records[0]["timestamp"] / 1e6
 
+    def get_agents(self, frame: int) -> List:
+        return [self._load_ego(frame=frame)]
+
     def _load_lidar(
         self,
         frame,

avapi/ugv/__init__.py

@@ -0,0 +1,4 @@
+from .dataset import UgvScenesManager
+
+
+__all__ = ["UgvScenesManager"]