fix typing

apperception-db · Aug 10, 2023 · 0d34c90 · 0d34c90
1 parent 12619ac
commit 0d34c90
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 7 deletions.
diff --git a/spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py b/spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py
@@ -114,7 +114,7 @@ class SegmentPoint:
     timestamp: "datetime.datetime"
     segment_type: "str"
     segment_line: "postgis.LineString"
-    segment_heading: "postgis.Polygon"
+    segment_heading: "float"
     road_polygon_info: "RoadPolygonInfo"
     obj_id: "int | None" = None
     type: "str | None" = None

diff --git a/spatialyze/video_processor/utils/associate_segment_mapping.py b/spatialyze/video_processor/utils/associate_segment_mapping.py
@@ -8,4 +8,4 @@ def associate_segment_mapping(
 ):
     if segment_mapping_meta is None:
         return None
-    return segment_mapping_meta[tracking_result.frame_idx].get(tracking_result.object_id)
+    return segment_mapping_meta[tracking_result.frame_idx].get(int(tracking_result.object_id))
diff --git a/spatialyze/video_processor/utils/format_trajectory.py b/spatialyze/video_processor/utils/format_trajectory.py
@@ -1,3 +1,7 @@
+import datetime
+import numpy.typing as npt
+import numpy as np
+
 from ..camera_config import CameraConfig
 from ..stages.segment_trajectory.construct_segment_trajectory import SegmentPoint
 from ..stages.tracking_3d.tracking_3d import Tracking3DResult
@@ -10,9 +14,9 @@ def format_trajectory(
     track: "list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]",
     base=None,
 ):
-    timestamps: "list[str]" = []
-    pairs: "list[Float3]" = []
-    itemHeadings: "list[int]" = []
+    timestamps: "list[datetime.datetime]" = []
+    pairs: "list[npt.NDArray[np.floating]]" = []
+    itemHeadings: "list[float | None]" = []
     translations: "list[Float3]" = []
     camera_id = None
     object_type = None
@@ -25,6 +29,7 @@ def format_trajectory(
     for tracking_result_3d, ego_info, segment_mapping in track:
         if ego_info:
             if (
+                ego_info.filename is not None and
                 "sweeps/CAM_FRONT/n008-2018-08-30-15-16-55-0400__CAM_FRONT__1535657125362404.jpg"
                 in ego_info.filename
             ):

diff --git a/spatialyze/video_processor/utils/get_tracks.py b/spatialyze/video_processor/utils/get_tracks.py
@@ -11,8 +11,8 @@ def get_tracks(
     ego_meta: "list[CameraConfig]",
     segment_mapping_meta: "list[SegmentTrajectoryMetadatum] | None" = None,
     base=None,
-) -> "dict[str, list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]]":
-    trajectories: "dict[str, list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]]" = {}
+) -> "dict[int, list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]]":
+    trajectories: "dict[int, list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]]" = {}
     for i in range(len(sortmeta)):
         frame = sortmeta[i]
         for obj_id, tracking_result in frame.items():