Add kwargs many places for kitti and nuscenes compat

Author: roshambo919

avapi/_dataset.py

@@ -34,6 +34,15 @@ def wrap_minus_pi_to_pi(phases):
     return phases
 
 
+def get_reference_from_line(line):
+    """A bit of a hack"""
+    line = line.split()
+    assert line[0] == "origin", line[0]
+    x = np.array([float(l) for l in line[1:4]])
+    q = np.quaternion(*[float(l) for l in line[4:8]])
+    return ReferenceFrame(x=x, q=q, reference=GlobalOrigin3D)
+
+
 class BaseSceneManager:
     def __iter__(self):
         for scene in self.scenes:
@@ -135,7 +144,7 @@ def get_sensor_name(self, sensor, agent=None):
 
     def get_frames(self, sensor, agent=None):
         sensor = self.get_sensor_name(sensor, agent=agent)
-        return self._load_frames(sensor, agent=agent)
+        return self._load_frames(sensor=sensor, agent=agent)
 
     def get_calibration(self, frame, sensor, agent=None):
         sensor = self.get_sensor_name(sensor, agent)
@@ -414,7 +423,7 @@ def read_dict_text_file(filepath):
     def parse_label_line(self, label_file_line):
         # Parse data elements
         data = label_file_line.strip("\n").split(" ")
-        if data[0] in ["avstack", "nuscenes"]:
+        if data[0] == "nuscenes":
             idx = 2
             ts = data[idx]
             idx += 1
@@ -669,10 +678,10 @@ def _get_sensor_file_name(self, frame, sensor=None, agent=None):
             ],
         )
 
-    def _load_frames(self, sensor: str = None):
+    def _load_frames(self, **kwargs):
         return self.frames
 
-    def _load_calibration(self, frame, ego_reference, sensor=None):
+    def _load_calibration(self, frame, ego_reference, sensor=None, **kwargs):
         """
         W := global "world" frame
         E := ego frame
@@ -699,18 +708,20 @@ def _load_calibration(self, frame, ego_reference, sensor=None):
             calib = calibration.Calibration(reference)
         return calib
 
-    def _load_timestamp(self, frame, sensor, utime=False):
+    def _load_timestamp(self, frame, sensor, utime=False, **kwargs):
+        if sensor is None:
+            sensor = "LIDAR_TOP"
         ut = self._get_sensor_record(frame, sensor)["timestamp"]
         if utime:
             return ut
         else:
             return ut / 1e6 - self.t0
 
-    def _load_image(self, frame, sensor=None):
+    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):
+    def _load_ego(self, frame, **kwargs):
         ref = GlobalOrigin3D
         if self.vehicle_pose is not None:
             try:
@@ -769,6 +780,7 @@ def _load_objects(
         sensor=None,
         whitelist_types=["car", "pedestrian", "bicycle", "truck", "bus", "motorcycle"],
         ignore_types=[],
+        **kwargs,
     ):
         """
         automatically loads into local sensor coordinates

avapi/kitti/dataset.py

@@ -1,12 +1,3 @@
-# -*- coding: utf-8 -*-
-# @Author: Spencer H
-# @Date:   2022-08-08
-# @Last Modified by:   Spencer H
-# @Last Modified date: 2022-09-29
-# @Description:
-"""
-
-"""
 import glob
 import os
 import re
@@ -156,10 +147,10 @@ def frames(self):
         """Wrapper to the classmethod using the split path"""
         return self._get_frames_folder(self.split_path)
 
-    def _load_frames(self, sensor: str):
+    def _load_frames(self, **kwargs):
         return self.frames
 
-    def _load_ego(self, frame):
+    def _load_ego(self, frame: int, **kwargs):
         """Ego origin is defined as the center of the vehicle"""
         reference = self._load_ego_reference(frame)
         ego = VehicleState("car")
@@ -215,7 +206,7 @@ def write_calibration(calib_dict, path, idx):
     def _load_ego_reference(self, frame):
         return GlobalOrigin3D  # TODO: this is not ideal...
 
-    def _load_calibration(self, frame, sensor, ego_reference):
+    def _load_calibration(self, frame: int, sensor: str, ego_reference, **kwargs):
         calib_fname = os.path.join(
             self.split_path, self.folder_names["calibration"], "%06d.txt" % frame
         )
@@ -277,7 +268,7 @@ def _load_calibration_from_file(
             raise NotImplementedError(sensor)
         return calib
 
-    def _load_image(self, frame, sensor):
+    def _load_image(self, frame: int, sensor: str, **kwargs):
         if isinstance(sensor, str) and "image" in sensor:
             img_fname = os.path.join(
                 self.split_path, self.folder_names[sensor], "%06d.png" % frame
@@ -290,8 +281,7 @@ def _load_image(self, frame, sensor):
             )
         return imread(img_fname)
 
-    def _load_lidar(self, frame, sensor, **kwargs):
-        filter_front = True  # always filter KITTI to front-only
+    def _load_lidar(self, frame: int, filter_front=True, **kwargs):
         lidar_fname = os.path.join(
             self.split_path, self.folder_names["lidar"], "%06d.bin" % frame
         )
@@ -309,6 +299,7 @@ def _load_objects(
         sensor="image-2",
         whitelist_types=["Car", "Pedestrian", "Cyclist"],
         ignore_types=["DontCare"],
+        **kwargs,
     ):
         label_fname = os.path.join(
             self.split_path, self.folder_names["label"], "%06d.txt" % frame
@@ -321,7 +312,7 @@ def _load_objects(
             obj.change_reference(object_calib.reference, inplace=True)
         return objects
 
-    def _load_timestamp(self, frame, utime=False, sensor="lidar"):
+    def _load_timestamp(self, frame: int, sensor: str = "lidar", **kwargs):
         if os.path.exists(
             os.path.join(self.split_path, self.folder_names["timestamps"])
         ):
@@ -349,7 +340,7 @@ def _save_objects(self, frame, objects, folder, file):
             for line in str_list:
                 f.write(f"{line}\n")
 
-    def _get_sensor_file_name(self, frame, sensor: str):
+    def _get_sensor_file_name(self, frame: int, sensor: str):
         return os.path.join(
             self.split_path,
             self.folder_names[sensor],

avapi/nuscenes/dataset.py

@@ -169,7 +169,12 @@ def make_sample_records(self):
         self.t0 = self.sample_records[0]["timestamp"] / 1e6
 
     def _load_lidar(
-        self, frame, sensor="LIDAR_TOP", filter_front=False, with_panoptic=False
+        self,
+        frame,
+        sensor="LIDAR_TOP",
+        filter_front=False,
+        with_panoptic=False,
+        **kwargs,
     ):
         if sensor.lower() == "lidar":
             sensor = self.sensors["lidar"]
@@ -187,7 +192,7 @@ def _load_lidar(
         else:
             return lidar
 
-    def _load_radar(self, frame, sensor="RADAR_FRONT"):
+    def _load_radar(self, frame, sensor="RADAR_FRONT", **kwargs):
         """ "
         Important: when trying to incorporate the velocity component of
         the radar data, make sure to look at the nuScenes forums for

avapi/utils.py

@@ -1,13 +1,3 @@
-# -*- coding: utf-8 -*-
-# @Author: Spencer H
-# @Date:   2022-04-03
-# @Last Modified by:   Spencer H
-# @Last Modified date: 2022-08-11
-# @Description:
-"""
-
-"""
-
 import os
 
 

avapi/visualize/movie.py

@@ -35,7 +35,7 @@ def make_movie_from_DM(
         i_end = len(DM.frames)
 
     imgs = []
-    boxes = []
+    boxes_out = []
     for frame_idx in range(i_start, i_end, 1):
         frame = DM.get_frames(sensor=CAM)[frame_idx]
         if len(boxes) == 0:
@@ -44,12 +44,12 @@ def make_movie_from_DM(
             img_boxes = boxes[frame_idx]
         img = DM.get_image(frame, sensor=CAM)
 
-        boxes.append(img_boxes)
+        boxes_out.append(img_boxes)
         imgs.append(img)
 
     make_movie(
         imgs,
-        boxes,
+        boxes_out,
         fps=DM.framerate,
         name=dataset_name,
         save=save,