Pin numpy to 1.21.2

argoverse/data_loading/argoverse_forecasting_loader.py

@@ -77,7 +77,7 @@ def seq_df(self) -> pd.DataFrame:
         return _read_csv(self.current_seq)
 
     @property
-    def agent_traj(self) -> np.ndarray:
+    def agent_traj(self) -> Any:
         """Get the trajectory for the track of type 'AGENT' in the current sequence.
 
         Returns:

argoverse/data_loading/argoverse_tracking_loader.py

@@ -6,6 +6,7 @@
 from typing import Dict, Iterator, List, Optional, Union, cast
 
 import numpy as np
+from numpy.typing import NDArray
 
 import argoverse.data_loading.object_label_record as object_label
 from argoverse.data_loading.object_label_record import ObjectLabelRecord
@@ -25,7 +26,7 @@ def __init__(self, root_dir: str) -> None:
         self.CAMERA_LIST = CAMERA_LIST
         self._log_list: Optional[List[str]] = None
         self._image_list: Optional[Dict[str, Dict[str, List[str]]]] = None
-        self._image_list_sync: Optional[Dict[str, Dict[str, List[np.ndarray]]]] = None
+        self._image_list_sync: Optional[Dict[str, Dict[str, List[NDArray[np.float64]]]]] = None
         self._lidar_list: Optional[Dict[str, List[str]]] = None
         self._image_timestamp_list: Optional[Dict[str, Dict[str, List[int]]]] = None
         self._timestamp_image_dict: Optional[Dict[str, Dict[str, Dict[int, str]]]] = None
@@ -128,7 +129,7 @@ def image_list(self) -> Dict[str, List[str]]:
         return self._image_list[self.current_log]
 
     @property
-    def image_list_sync(self) -> Dict[str, List[np.ndarray]]:
+    def image_list_sync(self) -> Dict[str, List[NDArray[np.float64]]]:
         """return list of image path (str) for all cameras for the current log.
 
         The different between image_list and image_list_sync is that image_list_sync
@@ -429,7 +430,7 @@ def get_image_at_timestamp(
         camera: str,
         log_id: Optional[str] = None,
         load: bool = True,
-    ) -> Optional[Union[str, np.ndarray]]:
+    ) -> Optional[Union[str, NDArray[np.float64]]]:
         """get image or image path at a specific timestamp
 
         Args:
@@ -459,7 +460,7 @@ def get_image_at_timestamp(
 
     def get_image(
         self, idx: int, camera: str, log_id: Optional[str] = None, load: bool = True
-    ) -> Union[str, np.ndarray]:
+    ) -> Union[str, NDArray[np.float64]]:
         """get image or image path at a specific index (in image index)
 
         Args:
@@ -488,7 +489,7 @@ def get_image(
 
     def get_image_sync(
         self, idx: int, camera: str, log_id: Optional[str] = None, load: bool = True
-    ) -> Union[str, np.ndarray]:
+    ) -> Union[str, NDArray[np.float64]]:
         """get image or image path at a specific index (in lidar index)
 
         Args:
@@ -515,7 +516,7 @@ def get_image_sync(
             return load_image(image_path)
         return image_path
 
-    def get_lidar(self, idx: int, log_id: Optional[str] = None, load: bool = True) -> Union[str, np.ndarray]:
+    def get_lidar(self, idx: int, log_id: Optional[str] = None, load: bool = True) -> Union[str, NDArray[np.float64]]:
         """Get lidar corresponding to frame index idx (in lidar frame).
 
         Args:

argoverse/data_loading/frame_label_accumulator.py

@@ -10,6 +10,7 @@
 from typing import Dict, List, Optional, Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 from argoverse.data_loading.frame_record import FrameRecord
 from argoverse.data_loading.object_label_record import ObjectLabelRecord
@@ -64,11 +65,11 @@ def __init__(
 
         # coordinate system is the map world frame
 
-        self.per_city_traj_dict: Dict[str, List[Tuple[np.ndarray, str]]] = {
+        self.per_city_traj_dict: Dict[str, List[Tuple[NDArray[np.float64], str]]] = {
             "MIA": [],
             "PIT": [],
         }  # all the trajectories for these 2 cities
-        self.log_egopose_dict: Dict[str, Dict[int, Dict[str, np.ndarray]]] = {}
+        self.log_egopose_dict: Dict[str, Dict[int, Dict[str, NDArray[np.float64]]]] = {}
         self.log_timestamp_dict: Dict[str, Dict[int, List[FrameRecord]]] = {}
         self.sdb = SynchronizationDB(self.dataset_dir)
 
@@ -146,7 +147,7 @@ def get_log_trajectory_labels(self, log_id: str) -> Optional[List[TrajectoryLabe
         else:
             return None
 
-    def place_trajectory_in_city_frame(self, traj_label: TrajectoryLabel, log_id: str) -> np.ndarray:
+    def place_trajectory_in_city_frame(self, traj_label: TrajectoryLabel, log_id: str) -> NDArray[np.float64]:
         """Place trajectory in the city frame
         Args:
             traj_label (TrajectoryLabel): instance of the TrajectoryLabel class.
@@ -218,8 +219,8 @@ def convert_bbox_to_city_frame(
         lidar_timestamp_ns: int,
         dataset_dir: str,
         log_id: str,
-        bbox_ego_frame: np.ndarray,
-    ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
+        bbox_ego_frame: NDArray[np.float64],
+    ) -> Tuple[NDArray[np.float64], Dict[str, NDArray[np.float64]]]:
         """Convert bounding box to city frame.
         Args:
             lidar_timestamp_ns (int): Lidar timestamp.

argoverse/data_loading/frame_record.py

@@ -5,6 +5,7 @@
 from typing import Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 
 class FrameRecord:
@@ -15,8 +16,8 @@ class FrameRecord:
 
     def __init__(
         self,
-        bbox_city_fr: np.ndarray,
-        bbox_ego_frame: np.ndarray,
+        bbox_city_fr: NDArray[np.float64],
+        bbox_ego_frame: NDArray[np.float64],
         occlusion_val: int,
         color: Tuple[float, float, float],
         track_uuid: str,

argoverse/data_loading/object_label_record.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 from argoverse.utils.calibration import CameraConfig, proj_cam_to_uv
 from argoverse.utils.cv2_plotting_utils import add_text_cv2, draw_clipped_line_segment
@@ -45,8 +46,8 @@ class ObjectLabelRecord:
 
     def __init__(
         self,
-        quaternion: np.ndarray,
-        translation: np.ndarray,
+        quaternion: NDArray[np.float64],
+        translation: NDArray[np.float64],
         length: float,
         width: float,
         height: float,
@@ -77,7 +78,7 @@ def __init__(
         self.track_id = track_id
         self.score = score
 
-    def as_2d_bbox(self) -> np.ndarray:
+    def as_2d_bbox(self) -> Any:
         """Convert the object cuboid to a 2D bounding box, with vertices provided in the egovehicle's reference frame.
 
         Length is x, width is y, and z is height
@@ -102,7 +103,7 @@ def as_2d_bbox(self) -> np.ndarray:
         bbox_in_egovehicle_frame = egovehicle_SE3_object.transform_point_cloud(bbox_object_frame)
         return bbox_in_egovehicle_frame
 
-    def as_3d_bbox(self) -> np.ndarray:
+    def as_3d_bbox(self) -> Any:
         r"""Calculate the 8 bounding box corners (returned as points inside the egovehicle's frame).
 
         Returns:
@@ -126,9 +127,9 @@ def as_3d_bbox(self) -> np.ndarray:
         The last four are the ones facing backwards.
         """
         # 3D bounding box corners. (Convention: x points forward, y to the left, z up.)
-        x_corners = self.length / 2 * np.array([1, 1, 1, 1, -1, -1, -1, -1])
-        y_corners = self.width / 2 * np.array([1, -1, -1, 1, 1, -1, -1, 1])
-        z_corners = self.height / 2 * np.array([1, 1, -1, -1, 1, 1, -1, -1])
+        x_corners: NDArray[np.float64] = self.length / 2 * np.array([1, 1, 1, 1, -1, -1, -1, -1])
+        y_corners: NDArray[np.float64] = self.width / 2 * np.array([1, -1, -1, 1, 1, -1, -1, 1])
+        z_corners: NDArray[np.float64] = self.height / 2 * np.array([1, 1, -1, -1, 1, 1, -1, -1])
         corners_object_frame = np.vstack((x_corners, y_corners, z_corners)).T
 
         egovehicle_SE3_object = SE3(rotation=quat2rotmat(self.quaternion), translation=self.translation)
@@ -137,17 +138,21 @@ def as_3d_bbox(self) -> np.ndarray:
 
     def render_clip_frustum_cv2(
         self,
-        img: np.ndarray,
-        corners: np.ndarray,
-        planes: List[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]],
+        img: NDArray[np.float64],
+        corners: NDArray[np.float64],
+        planes: List[
+            Tuple[
+                NDArray[np.float64], NDArray[np.float64], NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]
+            ]
+        ],
         camera_config: CameraConfig,
         colors: Tuple[Tuple[int, int, int], Tuple[int, int, int], Tuple[int, int, int]] = (
             BLUE_RGB,
             RED_RGB,
             GREEN_RGB,
         ),
         linewidth: int = 2,
-    ) -> np.ndarray:
+    ) -> NDArray[np.float64]:
         r"""We bring the 3D points into each camera, and do the clipping there.
 
         Renders box using OpenCV2. Edge coloring and vertex ordering is roughly based on
@@ -182,7 +187,7 @@ def render_clip_frustum_cv2(
             img: Numpy array of shape (M,N,3), representing updated image
         """
 
-        def draw_rect(selected_corners: np.ndarray, color: Tuple[int, int, int]) -> None:
+        def draw_rect(selected_corners: NDArray[np.float64], color: Tuple[int, int, int]) -> None:
             prev = selected_corners[-1]
             for corner in selected_corners:
                 draw_clipped_line_segment(
@@ -241,34 +246,34 @@ def draw_rect(selected_corners: np.ndarray, color: Tuple[int, int, int]) -> None
         return img
 
 
-def uv_coord_is_valid(uv: np.ndarray, img: np.ndarray) -> bool:
+def uv_coord_is_valid(uv: NDArray[np.float64], img: NDArray[np.float64]) -> bool:
     """Check if 2d-point lies within 3-channel color image boundaries"""
     h, w, _ = img.shape
     return bool(uv[0] >= 0 and uv[1] >= 0 and uv[0] < w and uv[1] < h)
 
 
-def label_is_closeby(box_point: np.ndarray) -> bool:
+def label_is_closeby(box_point: NDArray[np.float64]) -> bool:
     """Check if 3d cuboid pt (in egovehicle frame) is within range from
     egovehicle to prevent plot overcrowding.
     """
     return bool(np.linalg.norm(box_point) < MAX_RANGE_THRESH_PLOT_CATEGORY)
 
 
 def draw_alpha_rectangle(
-    img: np.ndarray,
+    img: NDArray[np.float64],
     top_left: Tuple[int, int],
     bottom_right: Tuple[int, int],
     color_rgb: Tuple[int, int, int],
     alpha: float,
-) -> np.ndarray:
+) -> NDArray[np.float64]:
     """Alpha blend colored rectangle into image. Corner coords given as (x,y) tuples"""
     img_h, img_w, _ = img.shape
     mask = np.zeros((img_h, img_w), dtype=np.uint8)
     mask[top_left[1] : bottom_right[1], top_left[0] : bottom_right[0]] = 1
     return vis_mask(img, mask, np.array(list(color_rgb[::-1])), alpha)
 
 
-def form_obj_label_from_json(label: Dict[str, Any]) -> Tuple[np.ndarray, str]:
+def form_obj_label_from_json(label: Dict[str, Any]) -> Tuple[NDArray[np.float64], str]:
     """Construct object from loaded json.
 
      The dictionary loaded from saved json file is expected to have the

argoverse/data_loading/stereo_dataloader.py

@@ -5,6 +5,7 @@
 
 import cv2
 import numpy as np
+from numpy.typing import NDArray
 
 from argoverse.utils.json_utils import read_json_file
 
@@ -76,7 +77,7 @@ def get_ordered_log_disparity_map_fpaths(self, log_id: str, disparity_name: str)
 
         return disparity_map_fpaths
 
-    def get_rectified_stereo_image(self, stereo_img_path: str) -> np.ndarray:
+    def get_rectified_stereo_image(self, stereo_img_path: str) -> Any:
         """Get the rectified stereo image.
 
         Args:
@@ -87,7 +88,7 @@ def get_rectified_stereo_image(self, stereo_img_path: str) -> np.ndarray:
         """
         return cv2.cvtColor(cv2.imread(stereo_img_path), cv2.COLOR_BGR2RGB)
 
-    def get_disparity_map(self, disparity_map_path: str) -> np.ndarray:
+    def get_disparity_map(self, disparity_map_path: str) -> NDArray[np.float64]:
         """Get the disparity map.
 
         The disparity maps are saved as uint16 PNG images. A zero-value ("0") indicates that no ground truth exists

argoverse/data_loading/synchronization_database.py

@@ -8,6 +8,7 @@
 from typing import Dict, Iterable, Optional, Tuple, cast
 
 import numpy as np
+from numpy.typing import NDArray
 from typing_extensions import Final
 
 from argoverse.sensor_dataset_config import ArgoverseConfig
@@ -38,7 +39,7 @@
 LIDAR_SWEEP_INTERVAL_W_BUFFER_MS = LIDAR_SWEEP_INTERVAL_MS + ALLOWED_TIMESTAMP_BUFFER_MS
 
 
-def get_timestamps_from_sensor_folder(sensor_folder_wildcard: str) -> np.ndarray:
+def get_timestamps_from_sensor_folder(sensor_folder_wildcard: str) -> NDArray[np.float64]:
     """Timestamp always lies at end of filename
 
     Args:
@@ -55,7 +56,7 @@ def get_timestamps_from_sensor_folder(sensor_folder_wildcard: str) -> np.ndarray
     return np.array([int(Path(jpg_fpath).stem.split("_")[-1]) for jpg_fpath in path_generator])
 
 
-def find_closest_integer_in_ref_arr(query_int: int, ref_arr: np.ndarray) -> Tuple[int, int]:
+def find_closest_integer_in_ref_arr(query_int: int, ref_arr: NDArray[np.float64]) -> Tuple[int, int]:
     """
     Find the closest integer to any integer inside a reference array, and the corresponding
     difference.
@@ -121,8 +122,8 @@ def __init__(self, dataset_dir: str, collect_single_log_id: Optional[str] = None
         else:
             log_fpaths = [f"{dataset_dir}/{collect_single_log_id}"]
 
-        self.per_log_camtimestamps_index: Dict[str, Dict[str, np.ndarray]] = {}
-        self.per_log_lidartimestamps_index: Dict[str, np.ndarray] = {}
+        self.per_log_camtimestamps_index: Dict[str, Dict[str, NDArray[np.float64]]] = {}
+        self.per_log_lidartimestamps_index: Dict[str, NDArray[np.float64]] = {}
 
         for log_fpath in log_fpaths:
             log_id = Path(log_fpath).name

argoverse/data_loading/trajectory_loader.py

@@ -10,6 +10,7 @@
 from typing import List, NamedTuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 from argoverse.data_loading.object_classes import OBJ_CLASS_MAPPING_DICT
 
@@ -36,20 +37,20 @@ class TrajectoryLabel(NamedTuple):
         heights (np.array): Array of heights for trajectory.
     """
 
-    timestamps: np.ndarray
-    quaternions: np.ndarray
-    translations: np.ndarray
+    timestamps: NDArray[np.float64]
+    quaternions: NDArray[np.float64]
+    translations: NDArray[np.float64]
     obj_class: int
     obj_class_str: str
-    occlusion: np.ndarray
+    occlusion: NDArray[np.float64]
     track_uuid: str
     log_id: str
     max_length: float
     max_width: float
     max_height: float
-    lengths: np.ndarray
-    widths: np.ndarray
-    heights: np.ndarray
+    lengths: NDArray[np.float64]
+    widths: NDArray[np.float64]
+    heights: NDArray[np.float64]
 
 
 def load_json_track_labels(log_track_labels_dir: str) -> List[TrajectoryLabel]:

argoverse/data_loading/vector_map_loader.py

@@ -45,6 +45,7 @@
 from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Tuple, Union, cast
 
 import numpy as np
+from numpy.typing import NDArray
 
 from argoverse.map_representation.lane_segment import LaneSegment
 
@@ -189,7 +190,7 @@ def get_lane_identifier(child: ET.Element) -> int:
     return int(child.attrib["lane_id"])
 
 
-def convert_node_id_list_to_xy(node_id_list: List[int], all_graph_nodes: Mapping[int, Node]) -> np.ndarray:
+def convert_node_id_list_to_xy(node_id_list: List[int], all_graph_nodes: Mapping[int, Node]) -> NDArray[np.float64]:
     """
     convert node id list to centerline xy coordinate