From d910bfe54c37c8bd86e2258021eb9d4c9f605c32 Mon Sep 17 00:00:00 2001 From: "Chanwut (Mick) Kittivorawong" <30903997+chanwutk@users.noreply.github.com> Date: Sun, 13 Aug 2023 19:47:26 -0700 Subject: [PATCH] Fix InView Typing (#17) * update tracking3d * style: [CI] format * update ablation * style: [CI] format * update type * fix type * fix types * do not change environment name * fix typing * fix type * style: [CI] format * clean up * fix lint * fix type error * style: [CI] format * fix typo * update * style: [CI] format * . * style: [CI] format --------- Co-authored-by: Github Actions Bot --- playground/run-ablation.ipynb | 24 +- playground/run-all-local | 10 + pyproject.toml | 4 +- .../stages/detection_2d/object_type_filter.py | 3 +- .../detection_estimation/segment_mapping.py | 4 +- .../stages/detection_estimation/utils.py | 3 +- .../video_processor/stages/in_view/in_view.py | 69 ++-- .../stages/segment_trajectory/__init__.py | 26 +- .../construct_segment_trajectory.py | 52 ++- .../segment_trajectory/from_tracking_3d.py | 24 +- .../stages/segment_trajectory/test.py | 365 ------------------ .../tracking_3d/from_tracking_2d_and_depth.py | 2 +- .../from_tracking_2d_and_detection_3d.py | 73 ++++ .../tracking_3d/from_tracking_2d_and_road.py | 13 +- .../stages/tracking_3d/tracking_3d.py | 13 +- .../utils/format_trajectory.py | 28 +- .../video_processor/utils/get_tracks.py | 15 +- .../video_processor/utils/infer_heading.py | 2 +- .../utils/insert_trajectory.py | 7 +- spatialyze/world.py | 51 +-- tests/workflow/test_optimized_workflow.py | 2 +- tests/workflow/test_simple_workflow.py | 2 +- 22 files changed, 276 insertions(+), 516 deletions(-) create mode 100644 playground/run-all-local delete mode 100644 spatialyze/video_processor/stages/segment_trajectory/test.py create mode 100644 spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_detection_3d.py diff --git a/playground/run-ablation.ipynb b/playground/run-ablation.ipynb index 6e1b17b..2a29aaa 100644 --- a/playground/run-ablation.ipynb +++ b/playground/run-ablation.ipynb @@ -172,6 +172,7 @@ "from spatialyze.video_processor.stages.tracking_3d.from_tracking_2d_and_road import FromTracking2DAndRoad\n", "from spatialyze.video_processor.stages.tracking_3d.from_tracking_2d_and_depth import FromTracking2DAndDepth\n", "from spatialyze.video_processor.stages.tracking_3d.tracking_3d import Tracking3DResult, Tracking3D\n", + "from spatialyze.video_processor.stages.tracking_3d.from_tracking_2d_and_detection_3d import FromTracking2DAndDetection3D as FromT2DAndD3D\n", "\n", "from spatialyze.video_processor.stages.segment_trajectory import SegmentTrajectory\n", "from spatialyze.video_processor.stages.segment_trajectory.construct_segment_trajectory import SegmentPoint\n", @@ -362,7 +363,7 @@ " # names = set(sampled_scenes)\n", " filtered_videos = [\n", " n for n in videos\n", - " if n[6:10] in names # and 'FRONT' in n # and n.endswith('FRONT')\n", + " if n[6:10] in names and 'FRONT' in n # and n.endswith('FRONT')\n", " ]\n", " N = len(filtered_videos)\n", " print('# of filtered videos:', N)\n", @@ -420,14 +421,14 @@ "\n", " # Ingest Trackings\n", " ego_meta = frames.interpolated_frames\n", - " sortmeta = FromTracking2DAndRoad.get(output)\n", + " sortmeta = Tracking3D.get(output)\n", " assert sortmeta is not None\n", " segment_trajectories = FromTracking3D.get(output)\n", " tracks = get_tracks(sortmeta, ego_meta, segment_trajectories)\n", " for obj_id, track in tracks.items():\n", " trajectory = format_trajectory(name, obj_id, track)\n", " if trajectory:\n", - " insert_trajectory(database, *trajectory)\n", + " insert_trajectory(database, *trajectory[0])\n", "\n", " # Ingest Camera\n", " accs: 'ACameraConfig' = []\n", @@ -569,9 +570,6 @@ "\n", " # Object Filter\n", " if object_filter:\n", - " # if isinstance(object_filter, bool):\n", - " # object_filter = ['car', 'truck']\n", - " # TODO: filter objects based on predicate\n", " pipeline.add_filter(ObjectTypeFilter(predicate=predicate))\n", "\n", " # 3D Detection\n", @@ -592,10 +590,11 @@ " cache=ss_cache,\n", " ))\n", "\n", - " if geo_depth:\n", - " pipeline.add_filter(FromTracking2DAndRoad())\n", - " else:\n", - " pipeline.add_filter(FromTracking3DAndDepth())\n", + " pipeline.add_filter(FromT2DAndD3D())\n", + " # if geo_depth:\n", + " # pipeline.add_filter(FromTracking2DAndRoad())\n", + " # else:\n", + " # pipeline.add_filter(FromTracking2DAndDepth())\n", "\n", " # Segment Trajectory\n", " # pipeline.add_filter(FromTracking3D())\n", @@ -902,9 +901,8 @@ }, "outputs": [], "source": [ - "tests = ['de', 'optde', 'noopt', 'inview', 'objectfilter', 'geo', 'opt']\n", - "# tests = ['de', 'optde']\n", - "# tests = ['de']\n", + "tests = ['optde', 'de', 'noopt', 'inview', 'objectfilter', 'geo', 'opt']\n", + "tests = ['de', 'noopt', 'inview', 'objectfilter']\n", "# random.shuffle(tests)\n", "\n", "for _test in tests:\n", diff --git a/playground/run-all-local b/playground/run-all-local new file mode 100644 index 0000000..55878f3 --- /dev/null +++ b/playground/run-all-local @@ -0,0 +1,10 @@ +#!/bin/bash + +tmux new-session -d -s run-test -n run-test-local || echo 'hi' +tmux send-keys -t run-test-local 'docker container start mobilitydb' Enter +tmux send-keys -t run-test-local 'cd ~/Documents/spatialyze' Enter +tmux send-keys -t run-test-local 'rm -rf ./outputs/run/*' Enter +tmux send-keys -t run-test-local 'rm -rf ./run-ablation.py' Enter +tmux send-keys -t run-test-local 'python ./spatialyze/utils/ingest_road.py "./data/scenic/road-network/boston-seaport"' Enter +tmux send-keys -t run-test-local 'jupyter nbconvert --to python ./playground/run-ablation.ipynb && mv playground/run-ablation.py .' Enter +tmux send-keys -t run-test-local 'python run-ablation.py' Enter \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index d9f1f95..409c06e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,10 +105,10 @@ exclude = 'spatialyze/video_processor/modules' pythonVersion = '3.10' ignore = [ 'spatialyze/video_processor/modules', + 'spatialyze/video_processor/stages/detection_2d/ground_truth.py', 'spatialyze/video_processor/stages/detection_estimation', - 'spatialyze/video_processor/stages/segment_trajectory', + 'spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py', 'spatialyze/video_processor/utils/process_pipeline.py', 'spatialyze/video_processor/utils/query_analyzer.py', 'spatialyze/video_processor/utils/preprocess.py', - 'spatialyze/video_processor' ] diff --git a/spatialyze/video_processor/stages/detection_2d/object_type_filter.py b/spatialyze/video_processor/stages/detection_2d/object_type_filter.py index 57a5688..ac0dae6 100644 --- a/spatialyze/video_processor/stages/detection_2d/object_type_filter.py +++ b/spatialyze/video_processor/stages/detection_2d/object_type_filter.py @@ -23,10 +23,11 @@ class ObjectTypeFilter(Detection2D): def __init__(self, types: "list[str] | None" = None, predicate: "PredicateNode | None" = None): - assert (types is None) != (predicate is None), "Can only except either types or predicate" if types is None: + assert predicate is not None, "Can only except either types or predicate" self.types = list(FindType()(predicate)) else: + assert types is not None, "Can only except either types or predicate" self.types = types print("types", self.types) diff --git a/spatialyze/video_processor/stages/detection_estimation/segment_mapping.py b/spatialyze/video_processor/stages/detection_estimation/segment_mapping.py index 8225dc7..ebc4ff4 100644 --- a/spatialyze/video_processor/stages/detection_estimation/segment_mapping.py +++ b/spatialyze/video_processor/stages/detection_estimation/segment_mapping.py @@ -137,13 +137,13 @@ class RoadPolygonInfo: segment_headings: "list[float]" contains_ego: bool ego_config: "CameraConfig" - fov_lines: "Tuple[Float22, Float22]" + fov_lines: "tuple[Float22, Float22]" def __post_init__(self): if len(self.segment_lines) == 0: return - start_segment_map: "dict[Float2, Tuple[shapely.geometry.LineString, float]]" = {} + start_segment_map: "dict[Float2, tuple[shapely.geometry.LineString, float]]" = {} ends: "set[Float2]" = set() for line, heading in zip(self.segment_lines, self.segment_headings): start = line.coords[0] diff --git a/spatialyze/video_processor/stages/detection_estimation/utils.py b/spatialyze/video_processor/stages/detection_estimation/utils.py index d7309f4..a1311e9 100644 --- a/spatialyze/video_processor/stages/detection_estimation/utils.py +++ b/spatialyze/video_processor/stages/detection_estimation/utils.py @@ -12,7 +12,8 @@ if TYPE_CHECKING: from ...camera_config import CameraConfig from .detection_estimation import DetectionInfo - from .segment_mapping import CameraPolygonMapping, RoadPolygonInfo + from .optimized_segment_mapping import RoadPolygonInfo + from .segment_mapping import CameraPolygonMapping logger = logging.getLogger(__name__) diff --git a/spatialyze/video_processor/stages/in_view/in_view.py b/spatialyze/video_processor/stages/in_view/in_view.py index 20422bd..88e5d0b 100644 --- a/spatialyze/video_processor/stages/in_view/in_view.py +++ b/spatialyze/video_processor/stages/in_view/in_view.py @@ -48,16 +48,17 @@ def __init__( ): super().__init__() self.distance = distance - assert (roadtypes is None) != ( - predicate is None - ), "Can only except either segment_type or predicate" - self.roadtypes: "list[str] | None" = None - self.predicate: "PredicateNode | None" = None + assert ( + roadtypes is not None or predicate is not None + ), "At least one of roadtypes or predicate must be specified" if roadtypes is not None: + assert predicate is None, "Can only except either segment_type or predicate" self.roadtypes = roadtypes if isinstance(roadtypes, list) else [roadtypes] - if predicate is not None: + self.predicate = None + elif predicate is not None: + assert roadtypes is None, "Can only except either segment_type or predicate" self.roadtypes, self.predicate_str = create_inview_predicate(predicate) - self.predicate = eval(self.predicate_str) + self.predicate = eval(str(self.predicate_str)) def __repr__(self) -> str: return f"InView(distance={self.distance}, roadtype={self.roadtypes}, predicate={self.predicate_str})" @@ -190,7 +191,7 @@ def get_views(payload: "Payload", distance: "float" = 100, skip: "bool" = True): ) ) - _extrinsics = np.stack(extrinsics) + _extrinsics = np.stack(extrinsics, dtype=np.floating) assert _extrinsics.shape == (N, 3, 4), _extrinsics.shape # convert 4 corner points from pixel-coordinate to world-coordinate @@ -248,9 +249,13 @@ def visit_BoolOpNode(self, node: BoolOpNode): # print('annihilated', visited) return lit(annihilator) if all(isinstance(e, LiteralNode) for e in visited.exprs): - assert len({e.value for e in visited.exprs}) == 1, visited.exprs + exprs = visited.exprs + assert all(isinstance(e, LiteralNode) for e in exprs), exprs + assert len({e.value for e in exprs if isinstance(e, LiteralNode)}) == 1, visited.exprs # print('all', visited) - return lit(visited.exprs[0].value) + e = exprs[0] + assert isinstance(e, LiteralNode), e + return lit(e.value) exprs = [e for e in visited.exprs if not isinstance(e, LiteralNode)] return BoolOpNode(visited.op, exprs) if len(exprs) > 1 else exprs[0] @@ -294,7 +299,7 @@ def visit_CallNode(self, node: CallNode): else: assert isinstance(segment, CallNode), segment assert segment == F.same_region().fn, segment - assert len(segment.params) == 1, segment.param + assert len(segment.params) == 1, segment.params assert isinstance(segment.params[0], LiteralNode), segment.params[0] assert isinstance(segment.params[0].value, str), segment.params[0] return F.is_roadtype(segment.params[0]) @@ -405,20 +410,24 @@ def visit_BoolOpNode(self, node: BoolOpNode): return F.ignore_roadtype() # Boolean Absorption - is_roadtypes = [ - e.params[0].value.lower() - for e in _exprs - if isinstance(e, CallNode) and e.fn == IS_ROADTYPE + e_params = [e.params[0] for e in _exprs if isinstance(e, CallNode) and e.fn == IS_ROADTYPE] + assert all( + isinstance(e, LiteralNode) and isinstance(e.value, str) for e in e_params + ), e_params + _is_roadtypes = [ + v.lower() + for v in (e.value for e in e_params if isinstance(e, LiteralNode)) + if isinstance(v, str) ] nested_exprs = [e for e in _exprs if isinstance(e, BoolOpNode) and e.op != node.op] - assert len(is_roadtypes) + len(nested_exprs) == len(_exprs), ( - len(is_roadtypes), + assert len(_is_roadtypes) + len(nested_exprs) == len(_exprs), ( + len(_is_roadtypes), len(nested_exprs), len(_exprs), _exprs, ) - is_roadtypes: "set[str]" = set(is_roadtypes) + is_roadtypes: "set[str]" = set(_is_roadtypes) def is_absorbed(e: "CallNode | BoolOpNode"): if isinstance(e, CallNode): @@ -432,11 +441,25 @@ def is_absorbed(e: "CallNode | BoolOpNode"): else: assert e.op == node.op, (node.op, e.op) all_roadtype = all(isinstance(e, CallNode) and e.fn == IS_ROADTYPE for e in e.exprs) - return all_roadtype and {ee.params[0].value.lower() for ee in e.exprs}.issubset( - is_roadtypes - ) - - nested_exprs = [e for e in nested_exprs if not any(map(is_absorbed, e.exprs))] + if not all_roadtype: + return False + exprs = [ee.params[0] for ee in e.exprs if isinstance(ee, CallNode)] + assert all( + isinstance(e, LiteralNode) and isinstance(e.value, str) for e in exprs + ), exprs + return { + v.lower() + for v in (ee.value for ee in exprs if isinstance(ee, LiteralNode)) + if isinstance(v, str) + }.issubset(is_roadtypes) + + def all_not_absorbed(e: "BoolOpNode"): + exprs = e.exprs + assert all(isinstance(e, (CallNode, BoolOpNode)) for e in exprs), exprs + exprs = (e for e in exprs if isinstance(e, (CallNode, BoolOpNode))) + return not any(map(is_absorbed, exprs)) + + nested_exprs = filter(all_not_absorbed, nested_exprs) _exprs = [*map(F.is_roadtype, sorted(is_roadtypes)), *nested_exprs] if len(_exprs) == 1: diff --git a/spatialyze/video_processor/stages/segment_trajectory/__init__.py b/spatialyze/video_processor/stages/segment_trajectory/__init__.py index d8e4f02..885b39f 100644 --- a/spatialyze/video_processor/stages/segment_trajectory/__init__.py +++ b/spatialyze/video_processor/stages/segment_trajectory/__init__.py @@ -10,7 +10,13 @@ from ..detection_estimation.utils import trajectory_3d from ..stage import Stage from ..tracking_2d.strongsort import StrongSORT -from .construct_segment_trajectory import SegmentPoint, calibrate +from .construct_segment_trajectory import ( + InvalidSegmentPoint, + PolygonAndId, + SegmentPoint, + ValidSegmentPoint, + calibrate, +) SegmentTrajectoryMetadatum = Dict[int, SegmentPoint] @@ -51,7 +57,7 @@ def _run(self, payload: "Payload"): @classmethod def encode_json(cls, o: "Any"): - if isinstance(o, SegmentPoint): + if isinstance(o, ValidSegmentPoint): return { "detection_id": tuple(o.detection_id), "car_loc3d": o.car_loc3d, @@ -65,6 +71,16 @@ def encode_json(cls, o: "Any"): "next": None if o.next is None else tuple(o.next.detection_id), "prev": None if o.prev is None else tuple(o.prev.detection_id), } + if isinstance(o, InvalidSegmentPoint): + return { + "detection_id": tuple(o.detection_id), + "car_loc3d": o.car_loc3d, + "timestamp": str(o.timestamp), + "obj_id": o.obj_id, + "type": o.type, + "next": None if o.next is None else tuple(o.next.detection_id), + "prev": None if o.prev is None else tuple(o.prev.detection_id), + } if isinstance(o, (RoadPolygonInfo, RoadPolygonInfo_)): return { "id": o.id, @@ -78,6 +94,12 @@ def encode_json(cls, o: "Any"): "fov_lines": o.fov_lines, } + if isinstance(o, PolygonAndId): + return { + "id": o.id, + "polygon": str(o.polygon), + } + def construct_trajectory(source: "Payload"): obj_3d_trajectories: "dict[int, list[Tuple[trajectory_3d, DetectionInfo, int]]]" = {} diff --git a/spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py b/spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py index 198fd4e..ea64a20 100644 --- a/spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py +++ b/spatialyze/video_processor/stages/segment_trajectory/construct_segment_trajectory.py @@ -1,13 +1,14 @@ import datetime import math from dataclasses import dataclass -from typing import List, Tuple +from typing import List, NamedTuple, Tuple, Union import numpy as np import numpy.typing as npt import postgis import psycopg2 import psycopg2.sql +import shapely.geometry from plpygis import Geometry from shapely.geometry import Point from shapely.ops import nearest_points @@ -107,21 +108,40 @@ SegmentQueryResult = Tuple[str, postgis.Polygon, postgis.LineString, List[str], float] +class PolygonAndId(NamedTuple): + id: "str" + polygon: "shapely.geometry.Polygon" + + @dataclass -class SegmentPoint: +class ValidSegmentPoint: detection_id: "DetectionId" car_loc3d: "Float3" timestamp: "datetime.datetime" segment_type: "str" segment_line: "postgis.LineString" segment_heading: "float" - road_polygon_info: "RoadPolygonInfo" + road_polygon_info: "RoadPolygonInfo | PolygonAndId" obj_id: "int | None" = None type: "str | None" = None next: "SegmentPoint | None" = None prev: "SegmentPoint | None" = None +@dataclass +class InvalidSegmentPoint: + detection_id: "DetectionId" + car_loc3d: "Float3" + timestamp: "datetime.datetime" + obj_id: "int | None" = None + type: "str | None" = None + next: "SegmentPoint | None" = None + prev: "SegmentPoint | None" = None + + +SegmentPoint = Union[ValidSegmentPoint, InvalidSegmentPoint] + + def construct_new_road_segment_info(result: "list[SegmentQueryResult]"): """Construct new road segment info based on query result. @@ -163,7 +183,7 @@ def construct_new_road_segment_info(result: "list[SegmentQueryResult]"): def find_middle_segment( - current_segment: "SegmentPoint", next_segment: "SegmentPoint", payload: "Payload" + current_segment: "ValidSegmentPoint", next_segment: "ValidSegmentPoint", payload: "Payload" ): current_time = current_segment.timestamp next_time = next_segment.timestamp @@ -227,7 +247,7 @@ def find_middle_segment( new_segment_line, new_heading = get_segment_line(new_road_segment_info, intersection_center) timestamp = current_time + (next_time - current_time) / 2 - middle_segment = SegmentPoint( + middle_segment = ValidSegmentPoint( DetectionId(middle_idx, None), (*intersection_center, 0.0), timestamp, @@ -239,10 +259,10 @@ def find_middle_segment( def binary_search_segment( - current_segment: "SegmentPoint", - next_segment: "SegmentPoint", + current_segment: "ValidSegmentPoint", + next_segment: "ValidSegmentPoint", payload: "Payload", -) -> "list[SegmentPoint]": +) -> "list[ValidSegmentPoint]": if 0 <= next_segment.detection_id.frame_idx - current_segment.detection_id.frame_idx <= 1: # same detection or detections on consecutive frame return [] @@ -270,7 +290,7 @@ def interpolate_segment(args: "Tuple[int, CameraConfig]"): assert location.shape == (3,) assert location.dtype == np.dtype(np.float64) - return SegmentPoint( + return ValidSegmentPoint( DetectionId(index, None), tuple(location), frame.timestamp, @@ -308,8 +328,10 @@ def interpolate_segment(args: "Tuple[int, CameraConfig]"): return before + [middle_segment] + after -def complete_segment_trajectory(road_segment_trajectory: "list[SegmentPoint]", payload: "Payload"): - completed_segment_trajectory: "list[SegmentPoint]" = [road_segment_trajectory[0]] +def complete_segment_trajectory( + road_segment_trajectory: "list[ValidSegmentPoint]", payload: "Payload" +): + completed_segment_trajectory: "list[ValidSegmentPoint]" = [road_segment_trajectory[0]] for current_segment, next_segment in zip( road_segment_trajectory[:-1], road_segment_trajectory[1:] @@ -327,14 +349,14 @@ def calibrate( detection_infos: "list[DetectionInfo]", frame_indices: "list[int]", payload: "Payload", -) -> "list[SegmentPoint]": +) -> "list[ValidSegmentPoint]": """Calibrate the trajectory to the road segments. Given a trajectory and the corresponding detection infos, map the trajectory to the correct road segments. The returned value is a list of SegmentTrajectoryPoint. """ - road_segment_trajectory: "list[SegmentPoint]" = [] + road_segment_trajectory: "list[ValidSegmentPoint]" = [] for i in range(len(trajectory_3d)): current_point3d, timestamp = trajectory_3d[i] current_point = current_point3d[:2] @@ -357,7 +379,7 @@ def calibrate( if current_road_segment_heading is None or math.cos(math.radians(relative_heading)) > 0: road_segment_trajectory.append( - SegmentPoint( + ValidSegmentPoint( detection_info.detection_id, current_point3d, timestamp, @@ -411,7 +433,7 @@ def calibrate( new_road_segment_info = construct_new_road_segment_info(result) new_segment_line, new_heading = get_segment_line(current_road_segment_info, current_point3d) road_segment_trajectory.append( - SegmentPoint( + ValidSegmentPoint( detection_info.detection_id, current_point3d, timestamp, diff --git a/spatialyze/video_processor/stages/segment_trajectory/from_tracking_3d.py b/spatialyze/video_processor/stages/segment_trajectory/from_tracking_3d.py index 82cc9d1..2f14091 100644 --- a/spatialyze/video_processor/stages/segment_trajectory/from_tracking_3d.py +++ b/spatialyze/video_processor/stages/segment_trajectory/from_tracking_3d.py @@ -15,11 +15,15 @@ from ...types import DetectionId from ..detection_3d import Detection3D from ..detection_3d import Metadatum as Detection3DMetadatum -from ..detection_estimation.segment_mapping import RoadPolygonInfo from ..tracking.tracking import Metadatum as TrackingMetadatum from ..tracking.tracking import Tracking from . import SegmentTrajectory, SegmentTrajectoryMetadatum -from .construct_segment_trajectory import SegmentPoint +from .construct_segment_trajectory import ( + InvalidSegmentPoint, + PolygonAndId, + SegmentPoint, + ValidSegmentPoint, +) USEFUL_TYPES = ["lane", "lanegroup", "intersection"] @@ -50,7 +54,7 @@ def _run(self, payload: "Payload"): object_map: "dict[int, dict[DetectionId, torch.Tensor]]" = dict() for fidx, frame in enumerate(tracking): - for tracking_result in frame.values(): + for tracking_result in frame: did = tracking_result.detection_id oid = tracking_result.object_id if oid not in object_map: @@ -159,18 +163,12 @@ def invalid_segment_point( oid: "int", class_map: "list[str]", ): - return SegmentPoint( + return InvalidSegmentPoint( did, tuple(((det[6:9] + det[9:12]) / 2).tolist()), timestamp, - None, - None, - None, - None, oid, class_map[int(det[5])], - None, - None, ) @@ -186,7 +184,7 @@ def valid_segment_point( polygonid: "str", shapely_polygon: "shapely.geometry.Polygon", ): - return SegmentPoint( + return ValidSegmentPoint( did, tuple(((det[6:9] + det[9:12]) / 2).tolist()), timestamp, @@ -194,11 +192,9 @@ def valid_segment_point( segmentline, segmentheading, # A place-holder for Polygon that only contain polygon id and polygon - RoadPolygonInfo(polygonid, shapely_polygon, [], None, [], None, None, None), + PolygonAndId(polygonid, shapely_polygon), oid, class_map[int(det[5])], - None, - None, ) diff --git a/spatialyze/video_processor/stages/segment_trajectory/test.py b/spatialyze/video_processor/stages/segment_trajectory/test.py deleted file mode 100644 index 5c981c4..0000000 --- a/spatialyze/video_processor/stages/segment_trajectory/test.py +++ /dev/null @@ -1,365 +0,0 @@ -import datetime -import math - -from plpygis import Geometry - -from spatialyze.database import database - -from ..detection_estimation.detection_estimation import DetectionInfo -from ..detection_estimation.segment_mapping import RoadPolygonInfo -from .construct_segment_trajectory import calibrate, test_segment_query - - -def get_test_trajectory(test_trajectory_points): - """This function is for testing only.""" - start_time = datetime.datetime.now() - trajectory = [] - for i in range(len(test_trajectory_points)): - start_time += datetime.timedelta(seconds=5) - trajectory.append((test_trajectory_points[i], start_time)) - return trajectory - - -def get_test_detection_infos(test_trajectory, test_segments): - """This function is for testing only.""" - assert len(test_trajectory) == len(test_segments) - detection_infos = [] - for i in range(len(test_trajectory)): - point, timestamp = test_trajectory[i] - test_segments_for_current_point = test_segments[i] - road_segment_info = None - for test_segment in test_segments_for_current_point: - segmentid, segmentpolygon, segmentline, segmenttype, segmentheading = test_segment - segmentheading = math.degrees(segmentheading) if segmentheading is not None else None - segmentline = Geometry(segmentline.to_ewkb()).shapely if segmentline else None - if road_segment_info is not None: - road_segment_info.segment_lines.append(segmentline) - road_segment_info.segment_headings.append(segmentheading) - continue - segmentpolygon = Geometry(segmentpolygon.to_ewkb()).shapely - road_segment_info = RoadPolygonInfo( - segmentid, - segmentpolygon, - [segmentline], - segmenttype, - [segmentheading], - None, - False, - None, - ) - detection_info = DetectionInfo( - detection_id=segmentid, - frame_segment=None, - road_polygon_info=road_segment_info, - car_loc3d=point, - car_loc2d=None, - car_bbox3d=None, - car_bbox2d=None, - ego_trajectory=None, - ego_config=None, - ego_road_polygon_info=None, - timestamp=timestamp, - ) - detection_infos.append(detection_info) - return detection_infos - - -def test_same_segment(): - print("test same segment") - test_segment_ids = [ - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [(1955, 870), (1960, 874), (1980, 872), (1990, 875)] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "53f56897-4795-4d75-a721-3c969bb3206c", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -def test_wrong_start_same_segment(): - print("test wrong start same segment") - test_segment_ids = [ - "c67e592f-2e73-4165-b8cf-64165bb300a8", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [(1955, 874), (1960, 870), (1980, 872), (1990, 875)] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -def test_connected_segments(): - print("test connected segment") - """Some trajectory points are in the wrong segments.""" - test_segment_ids = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "e39e4059-3a55-42f9-896f-475d89a70e86", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "aa22ee59-c9ef-4759-a69c-c295469f3e37_inter", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [ - (1910, 869), - (1915, 873), - (1920, 871), - (1940, 870), - (1955, 870), - (1960, 874), - (1980, 872), - (1990, 875), - ] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "2e6d0881-bb10-4145-a45f-28382c46e476", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "aa22ee59-c9ef-4759-a69c-c295469f3e37_inter", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "53f56897-4795-4d75-a721-3c969bb3206c", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -def test_complete_story1(): - """Simplest complete story case. - - The trajectories are all in the correct segments. - """ - print("test complete story 1") - test_segment_ids = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [ - (1910, 869), - (1920, 871), - (1955, 870), - (1960, 871), - (1980, 871), - (1990, 871), - ] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "aa22ee59-c9ef-4759-a69c-c295469f3e37_inter", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -def test_complete_story2(): - """Some trajectory points are in the wrong segments.""" - print("test complete story 2") - test_segment_ids = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "e39e4059-3a55-42f9-896f-475d89a70e86", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "c67e592f-2e73-4165-b8cf-64165bb300a8", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [ - (1910, 869), - (1920, 874), - (1955, 870), - (1960, 874), - (1980, 872), - (1990, 875), - ] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "aa22ee59-c9ef-4759-a69c-c295469f3e37_inter", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - "99c90907-e7a2-4b19-becc-afe2b7f013c7", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -def test_complete_story3(): - """Some trajectory points are in the wrong segments.""" - print("test complete story 3") - test_segment_ids = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "e39e4059-3a55-42f9-896f-475d89a70e86", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "53c5901a-dad9-4f0d-bcb6-c127dda2be09", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "53c5901a-dad9-4f0d-bcb6-c127dda2be09", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [ - (1910, 869), - (1915, 873), - (1920, 871), - (1937, 882), - (1932, 885), - (1937, 887), - (1932, 892), - ] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "aa22ee59-c9ef-4759-a69c-c295469f3e37_inter", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -def test_complete_story4(): - """Most trajectory points are in the wrong segments.""" - print("test complete story 4") - test_segment_ids = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "e39e4059-3a55-42f9-896f-475d89a70e86", - "53c5901a-dad9-4f0d-bcb6-c127dda2be09", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "53c5901a-dad9-4f0d-bcb6-c127dda2be09", - ] - test_segments = [ - database.execute(test_segment_query.format(segment_id=segment_id)) - for segment_id in test_segment_ids - ] - test_trajectory_points = [ - (1910, 868), - (1920, 873), - (1932, 885), - (1937, 887), - (1932, 892), - ] - test_trajectory = get_test_trajectory(test_trajectory_points) - test_detection_infos = get_test_detection_infos(test_trajectory, test_segments) - segment_trajectory = calibrate(test_trajectory, test_detection_infos) - correct_result = [ - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "34c01bd5-f649-42e2-be32-30f9a4d02b25", - "aa22ee59-c9ef-4759-a69c-c295469f3e37_inter", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - "9eef6c56-c5d9-46ed-a44e-9848676bdddf", - ] - print( - [ - segment_trajectory_point.road_segment_info.segment_id - for segment_trajectory_point in segment_trajectory - ] - ) - print("correct result", correct_result) - - -if __name__ == "__main__": - test_same_segment() - test_wrong_start_same_segment() - test_connected_segments() - test_complete_story1() - test_complete_story2() - test_complete_story3() - test_complete_story4() - print("All tests passed!") diff --git a/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_depth.py b/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_depth.py index a663cfe..cc3eafa 100644 --- a/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_depth.py +++ b/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_depth.py @@ -43,7 +43,7 @@ def _run(self, payload: "Payload") -> "tuple[bitarray | None, dict[str, list] | t.detection_id, t.object_id, point_from_camera, - point, + tuple(point.tolist()), t.bbox_left, t.bbox_top, t.bbox_w, diff --git a/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_detection_3d.py b/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_detection_3d.py new file mode 100644 index 0000000..531f8f1 --- /dev/null +++ b/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_detection_3d.py @@ -0,0 +1,73 @@ +from bitarray import bitarray + +from ...payload import Payload +from ..detection_3d import Detection3D +from ..tracking_2d.tracking_2d import Tracking2D +from .tracking_3d import Metadatum, Tracking3D, Tracking3DResult + + +class FromTracking2DAndDetection3D(Tracking3D): + def _run(self, payload: "Payload") -> "tuple[bitarray | None, dict[str, list] | None]": + metadata: "list[Metadatum]" = [] + trajectories: "dict[int, list[Tracking3DResult]]" = {} + + trackings = Tracking2D.get(payload.metadata) + assert trackings is not None + + detections = Detection3D.get(payload.metadata) + assert detections is not None + + for k, detection, tracking, frame in zip( + payload.keep, detections, trackings, payload.video + ): + dets, _, dids = detection + if not k or tracking is None or detection is None: + metadata.append(dict()) + continue + + points_left = dets[:, 6:9] + points_right = dets[:, 9:12] + points = (points_left + points_right) / 2 + + points_from_camera_left = dets[:, 12:15] + points_from_camera_right = dets[:, 15:18] + points_from_camera = (points_from_camera_left + points_from_camera_right) / 2 + + detection_map = { + did: (det, p, pfc) + for det, p, pfc, did in zip( + dets, points.tolist(), points_from_camera.tolist(), dids + ) + } + trackings3d: "dict[int, Tracking3DResult]" = {} + for object_id, t in tracking.items(): + did = t.detection_id + det, p, pfc = detection_map[did] + + trackings3d[object_id] = Tracking3DResult( + t.frame_idx, + t.detection_id, + t.object_id, + tuple(pfc), + tuple(p), + t.bbox_left, + t.bbox_top, + t.bbox_w, + t.bbox_h, + t.object_type, + frame.timestamp, + ) + if object_id not in trajectories: + trajectories[object_id] = [] + trajectories[object_id].append(trackings3d[object_id]) + metadata.append(trackings3d) + + for trajectory in trajectories.values(): + last = len(trajectory) - 1 + for i, traj in enumerate(trajectory): + if i > 0: + traj.prev = trajectory[i - 1] + if i < last: + traj.next = trajectory[i + 1] + + return None, {self.classname(): metadata} diff --git a/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_road.py b/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_road.py index 34e6f40..9afc79e 100644 --- a/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_road.py +++ b/spatialyze/video_processor/stages/tracking_3d/from_tracking_2d_and_road.py @@ -63,21 +63,18 @@ def _run(self, payload: "Payload"): points = rotated_directions * ts + translation[:, np.newaxis] points_from_camera = rotate(points - translation[:, np.newaxis], rotation.inverse) - for t, oid, point, point_from_camera in zip(_ts, oids, points.T, points_from_camera.T): + for t, oid, point, point_from_camera in zip( + _ts, oids, points.T.tolist(), points_from_camera.T.tolist() + ): assert point_from_camera.shape == (3,) assert isinstance(oid, int) or oid.is_integer() oid = int(oid) - point_from_camera = ( - point_from_camera[0], - point_from_camera[1], - point_from_camera[2], - ) trackings3d[oid] = Tracking3DResult( t.frame_idx, t.detection_id, oid, - point_from_camera, - point, + tuple(point_from_camera), + tuple(point), t.bbox_left, t.bbox_top, t.bbox_w, diff --git a/spatialyze/video_processor/stages/tracking_3d/tracking_3d.py b/spatialyze/video_processor/stages/tracking_3d/tracking_3d.py index f6873c6..41b830b 100644 --- a/spatialyze/video_processor/stages/tracking_3d/tracking_3d.py +++ b/spatialyze/video_processor/stages/tracking_3d/tracking_3d.py @@ -1,11 +1,8 @@ import datetime from dataclasses import dataclass -from typing import Any, Dict, Tuple +from typing import Any, Dict -import numpy as np -import numpy.typing as npt - -from ...types import DetectionId +from ...types import DetectionId, Float3 from ..stage import Stage @@ -14,8 +11,8 @@ class Tracking3DResult: frame_idx: int detection_id: DetectionId object_id: float - point_from_camera: Tuple[float, float, float] - point: "npt.NDArray[np.floating]" + point_from_camera: "Float3" + point: "Float3" bbox_left: float bbox_top: float bbox_w: float @@ -38,7 +35,7 @@ def encode_json(cls, o: "Any"): "detection_id": tuple(o.detection_id), "object_id": o.object_id, "point_from_camera": o.point_from_camera, - "point": o.point.tolist(), + "point": o.point, "bbox_left": o.bbox_left, "bbox_top": o.bbox_top, "bbox_w": o.bbox_w, diff --git a/spatialyze/video_processor/utils/format_trajectory.py b/spatialyze/video_processor/utils/format_trajectory.py index de23007..8b25014 100644 --- a/spatialyze/video_processor/utils/format_trajectory.py +++ b/spatialyze/video_processor/utils/format_trajectory.py @@ -1,22 +1,13 @@ import datetime -import numpy as np -import numpy.typing as npt - -from ..camera_config import CameraConfig -from ..stages.segment_trajectory.construct_segment_trajectory import SegmentPoint -from ..stages.tracking_3d.tracking_3d import Tracking3DResult +from ..stages.segment_trajectory.construct_segment_trajectory import InvalidSegmentPoint from ..types import Float3 +from .get_tracks import TrackPoint -def format_trajectory( - video_name: "str", - obj_id: "int", - track: "list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]", - base=None, -): +def format_trajectory(video_name: "str", obj_id: "int", track: "list[TrackPoint]"): timestamps: "list[datetime.datetime]" = [] - pairs: "list[npt.NDArray[np.floating]]" = [] + pairs: "list[Float3]" = [] itemHeadings: "list[float | None]" = [] translations: "list[Float3]" = [] camera_id = None @@ -47,10 +38,13 @@ def format_trajectory( object_type = tracking_result_3d.object_type timestamps.append(ego_info.timestamp) pairs.append(tracking_result_3d.point) - if not segment_mapping or (segment_mapping.segment_type == "intersection"): - itemHeadings.append(None) - else: - itemHeadings.append(segment_mapping.segment_heading) + itemHeadings.append( + None + if segment_mapping is None + or isinstance(segment_mapping, InvalidSegmentPoint) + or segment_mapping.segment_type == "intersection" + else segment_mapping.segment_heading + ) translations.append(ego_info.ego_translation) # road_types.append(segment_mapping.road_polygon_info.road_type if base else detection_info.road_type) # roadpolygons.append(None if base else detection_info.road_polygon_info.polygon) diff --git a/spatialyze/video_processor/utils/get_tracks.py b/spatialyze/video_processor/utils/get_tracks.py index c496a64..2747141 100644 --- a/spatialyze/video_processor/utils/get_tracks.py +++ b/spatialyze/video_processor/utils/get_tracks.py @@ -1,3 +1,5 @@ +from typing import NamedTuple + from ..camera_config import CameraConfig from ..stages.segment_trajectory.construct_segment_trajectory import SegmentPoint from ..stages.segment_trajectory.from_tracking_3d import SegmentTrajectoryMetadatum @@ -6,13 +8,18 @@ from .associate_segment_mapping import associate_segment_mapping +class TrackPoint(NamedTuple): + tracking_result: "Tracking3DResult" + ego_info: "CameraConfig" + segment_mapping: "SegmentPoint | None" + + def get_tracks( sortmeta: "list[Tracking3DMetadatum]", ego_meta: "list[CameraConfig]", segment_mapping_meta: "list[SegmentTrajectoryMetadatum] | None" = None, - base=None, -) -> "dict[int, list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]]": - trajectories: "dict[int, list[tuple[Tracking3DResult, CameraConfig, SegmentPoint | None]]]" = {} +) -> "dict[int, list[TrackPoint]]": + trajectories: "dict[int, list[TrackPoint]]" = {} for i in range(len(sortmeta)): frame = sortmeta[i] for obj_id, tracking_result in frame.items(): @@ -23,7 +30,7 @@ def get_tracks( tracking_result, segment_mapping_meta ) trajectories[obj_id].append( - (tracking_result, associated_ego_info, associated_segment_mapping) + TrackPoint(tracking_result, associated_ego_info, associated_segment_mapping) ) for trajectory in trajectories.values(): diff --git a/spatialyze/video_processor/utils/infer_heading.py b/spatialyze/video_processor/utils/infer_heading.py index 6b6f239..aed25df 100644 --- a/spatialyze/video_processor/utils/infer_heading.py +++ b/spatialyze/video_processor/utils/infer_heading.py @@ -4,7 +4,7 @@ def infer_heading( - curItemHeading: "int | None", prevPoint: "Float3 | None", current_point: "Float3" + curItemHeading: "float | None", prevPoint: "Float3 | None", current_point: "Float3" ): if curItemHeading is not None: return math.degrees(curItemHeading) diff --git a/spatialyze/video_processor/utils/insert_trajectory.py b/spatialyze/video_processor/utils/insert_trajectory.py index c45285b..9ac99c3 100644 --- a/spatialyze/video_processor/utils/insert_trajectory.py +++ b/spatialyze/video_processor/utils/insert_trajectory.py @@ -1,8 +1,5 @@ import datetime -import numpy as np -import numpy.typing as npt - from ...database import Database from ...utils import join from ..types import Float3 @@ -15,7 +12,7 @@ def insert_trajectory( camera_id: str, object_type: str, postgres_timestamps: "list[datetime.datetime]", - pairs: "list[npt.NDArray[np.floating]]", + pairs: "list[Float3]", itemHeading_list: "list[float | None]", translation_list: "list[Float3]", # road_types: "list[str]", @@ -24,7 +21,7 @@ def insert_trajectory( traj_centroids: "list[str]" = [] translations: "list[str]" = [] itemHeadings: "list[str]" = [] - prevTimestamp: "str | None" = None + prevTimestamp: "datetime.datetime | None" = None prevPoint: "Float3 | None" = None for timestamp, current_point, curItemHeading, current_trans in zip( postgres_timestamps, pairs, itemHeading_list, translation_list diff --git a/spatialyze/world.py b/spatialyze/world.py index 91352e7..c3b309d 100644 --- a/spatialyze/world.py +++ b/spatialyze/world.py @@ -22,12 +22,10 @@ ) from .video_processor.stages.detection_estimation import DetectionEstimation from .video_processor.stages.in_view.in_view import InView +from .video_processor.stages.stage import Stage from .video_processor.stages.tracking_2d.strongsort import StrongSORT -from .video_processor.stages.tracking_3d.from_tracking_2d_and_depth import ( - FromTracking2DAndDepth, -) -from .video_processor.stages.tracking_3d.from_tracking_2d_and_road import ( - FromTracking2DAndRoad, +from .video_processor.stages.tracking_3d.from_tracking_2d_and_detection_3d import ( + FromTracking2DAndDetection3D, ) from .video_processor.stages.tracking_3d.tracking_3d import Metadatum as T3DMetadatum from .video_processor.stages.tracking_3d.tracking_3d import Tracking3D @@ -107,35 +105,24 @@ def _execute(world: "World", optimization=True): # for gc in world._geogConstructs: # gc.ingest(database) # analyze predicates to generate pipeline - objtypes_filter = ObjectTypeFilter(predicate=world.predicates) + steps: "list[Stage]" = [] + if optimization: + steps.append(InView(distance=50, predicate=world.predicates)) + steps.append(DecodeFrame()) + steps.append(YoloDetection()) if optimization: - pipeline = Pipeline( - [ - DecodeFrame(), - InView(distance=50, predicate=world.predicates), - YoloDetection(), - objtypes_filter, - FromDetection2DAndRoad(), - *( - [DetectionEstimation()] - if all(t in ["car", "truck"] for t in objtypes_filter.types) - else [] - ), - StrongSORT(), - FromTracking2DAndRoad(), - ] - ) + objtypes_filter = ObjectTypeFilter(predicate=world.predicates) + steps.append(objtypes_filter) + steps.append(FromDetection2DAndRoad()) + if all(t in ["car", "truck"] for t in objtypes_filter.types): + steps.append(DetectionEstimation()) else: - pipeline = Pipeline( - [ - DecodeFrame(), - YoloDetection(), - DepthEstimation(), - FromDetection2DAndDepth(), - StrongSORT(), - FromTracking2DAndDepth(), - ] - ) + steps.append(DepthEstimation()) + steps.append(FromDetection2DAndDepth()) + steps.append(StrongSORT()) + steps.append(FromTracking2DAndDetection3D()) + + pipeline = Pipeline(steps) qresults: "dict[str, list[tuple]]" = {} vresults: "dict[str, list[T3DMetadatum]]" = {} diff --git a/tests/workflow/test_optimized_workflow.py b/tests/workflow/test_optimized_workflow.py index 2904a9d..a483a9f 100644 --- a/tests/workflow/test_optimized_workflow.py +++ b/tests/workflow/test_optimized_workflow.py @@ -41,7 +41,7 @@ def test_optimized_workflow(): assert tuple(p.detection_id) == tuple(g.detection_id), (p.detection_id, g.detection_id) assert p.object_id == g.object_id, (p.object_id, g.object_id) assert np.allclose(np.array(p.point_from_camera), np.array(g.point_from_camera)), (p.point_from_camera, g.point_from_camera) - assert np.allclose(np.array(p.point.tolist()), np.array(g.point)), (p.point, g.point) + assert np.allclose(np.array(p.point), np.array(g.point)), (p.point, g.point) assert p.bbox_left == g.bbox_left, (p.bbox_left, g.bbox_left) assert p.bbox_top == g.bbox_top, (p.bbox_top, g.bbox_top) assert p.bbox_w == g.bbox_w, (p.bbox_w, g.bbox_w) diff --git a/tests/workflow/test_simple_workflow.py b/tests/workflow/test_simple_workflow.py index 8f51bcf..75c6ce6 100644 --- a/tests/workflow/test_simple_workflow.py +++ b/tests/workflow/test_simple_workflow.py @@ -41,7 +41,7 @@ def test_simple_workflow(): assert tuple(p.detection_id) == tuple(g.detection_id), (p.detection_id, g.detection_id) assert p.object_id == g.object_id, (p.object_id, g.object_id) assert np.allclose(np.array(p.point_from_camera), np.array(g.point_from_camera)), (p.point_from_camera, g.point_from_camera) - assert np.allclose(np.array(p.point.tolist()), np.array(g.point)), (p.point, g.point) + assert np.allclose(np.array(p.point), np.array(g.point)), (p.point, g.point) assert p.bbox_left == g.bbox_left, (p.bbox_left, g.bbox_left) assert p.bbox_top == g.bbox_top, (p.bbox_top, g.bbox_top) assert p.bbox_w == g.bbox_w, (p.bbox_w, g.bbox_w)