Add agent getting to nuscenes

Author: roshambo919

avapi/_dataset.py

@@ -134,7 +134,7 @@ def check_frame(self, frame):
     def get_agents(self, frame: int) -> "DataContainer":
         return self._load_agents(frame)
 
-    def get_agent(self, frame: int, agent: int):
+    def get_agent(self, frame: int, agent: int) -> List[VehicleState]:
         agents = self.get_agents(frame)
         return [ag for ag in agents if ag.ID == agent][0]
 
@@ -158,6 +158,9 @@ def get_sensor_name(self, sensor, agent=None) -> str:
             return sensor
         else:
             return self.sensors[sensor]
+        
+    def get_sensor_names_by_type(self, sensor_type: str, agent=None) -> List[str]:
+        return self._load_sensor_names_by_type(sensor_type=sensor_type, agent=agent)
 
     def get_frames(self, sensor, agent=None) -> List[int]:
         sensor = self.get_sensor_name(sensor, agent=agent)
@@ -344,6 +347,9 @@ def save_objects(self, frame, objects, folder, file=None):
             os.makedirs(folder)
         self._save_objects(frame, objects, folder, file=file)
 
+    def _load_sensor_names_by_type(self, sensor_type, agent):
+        raise NotImplementedError
+
     def _load_agents(self, frame):
         raise NotImplementedError
 
@@ -371,8 +377,8 @@ def _load_lidar(self, frame, sensor, agent=None):
     def _load_objects(self, frame, sensor, agent=None):
         raise NotImplementedError
 
-    def _load_sensor_data_filepath(self, frame, sensor: str):
-        return self._get_sensor_file_name(frame, sensor)
+    def _load_sensor_data_filepath(self, frame, sensor: str, agent=None):
+        return self._get_sensor_file_name(frame, sensor, agent=agent)
 
     def _load_objects_from_file(
         self,
@@ -752,7 +758,7 @@ def _load_image(self, frame, sensor=None, **kwargs):
         img_fname = self._get_sensor_file_name(frame, sensor)
         return imread(img_fname)
 
-    def _load_ego(self, frame, **kwargs):
+    def _load_ego(self, frame, **kwargs) -> VehicleState:
         ref = GlobalOrigin3D
         if self.vehicle_pose is not None:
             try:

avapi/nuscenes/dataset.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 from avstack.config import DATASETS
+from avstack.environment.objects import VehicleState
 from avstack.geometry.transformations import matrix_cartesian_to_spherical
 from scipy.interpolate import interp1d
 
@@ -164,9 +165,19 @@ def make_sample_records(self):
             )
         self.t0 = self.sample_records[0]["timestamp"] / 1e6
 
-    def get_agents(self, frame: int) -> List:
+    def get_agents(self, frame: int) -> List[VehicleState]:
         return [self._load_ego(frame=frame)]
 
+    def _load_sensor_names_by_type(self, sensor_type, **kwargs):
+        if sensor_type.lower() == "camera":
+            return [sID for sID in self.sensor_IDs if "CAM" in sID]
+        elif sensor_type.lower() == "lidar":
+            return [sID for sID in self.sensor_IDs if "LIDAR" in sID]
+        elif sensor_type.lower() == "radar":
+            return [sID for sID in self.sensor_IDs if "RADAR" in sID]
+        else:
+            raise NotImplementedError(sensor_type)
+
     def _load_lidar(
         self,
         frame,