From c6747ac8e68de3334afbd40da28880f0df2d0d43 Mon Sep 17 00:00:00 2001 From: Ruben van de Ven Date: Fri, 31 Oct 2025 19:21:49 +0100 Subject: [PATCH] draw training lines --- trap/anomaly.py | 181 +++++++++++++++++++++++++++++++++++++++++ trap/base.py | 1 + trap/lines.py | 87 +++++++++++++++++++- trap/renderable.proto | 50 ++++++++++++ trap/stage.py | 121 ++++++++++++++++++++++------ trap/track_history.py | 183 ++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 593 insertions(+), 30 deletions(-) create mode 100644 trap/anomaly.py create mode 100644 trap/renderable.proto create mode 100644 trap/track_history.py diff --git a/trap/anomaly.py b/trap/anomaly.py new file mode 100644 index 0000000..5664271 --- /dev/null +++ b/trap/anomaly.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import logging +from typing import List + +import numpy as np +from trap.base import ProjectedTrack +from trap.lines import AppendableLine, Coordinate, DeltaT, ProceduralChain, RenderableLines, SrgbaColor, StaticLine + +logger = logging.getLogger('anomaly') + +def calc_anomaly(segments: List[DiffSegment], window: int = 3): + """Calculate anomaly score based on provided segments + considering a sliding window of the last n items + """ + + relevant_segments = segments[-window:] + scores = [s.avg_score() for s in relevant_segments] + s = list(filter(lambda x: x is not None,scores)) + + return np.average(s) + + +class DiffSegment(): + """ + A segment of a prediction track, that can be diffed + with a track. The track is continously updated. + If a new prediction comes in, the diff is marked as + finished. After which it is animated and added to the + Scenario's anomaly score. + """ + DRAW_DECAY_SPEED = 25 + POINT_INTERVAL = 4 + + def __init__(self, prediction: ProjectedTrack): + self.ptrack = prediction + self._last_diff_frame_idx: int = 0 + self.finished = False + + self.line = StaticLine() + self.points: List[Coordinate] = [] + self._drawn_points = [] + self._target_track = prediction + + self.score = 0 + + def finish(self): + self.finished = True + + def nr_of_passed_points(self) -> int: + if not self._last_diff_frame_idx: + return 0 + return self._last_diff_frame_idx - self.ptrack.frame_index + + # if isinstance(self.line, AppendableLine): + # return self.line.nr_of_passed_points() * self.POINT_INTERVAL + # else: + # return len(self.points) * self.POINT_INTERVAL + + def avg_score(self): + frames_passed = self.nr_of_passed_points() + if not frames_passed: + return None + else: + return self.score/frames_passed + + + # run on each track update received + def update_track(self, track: ProjectedTrack): + self._target_track = track + + if self.finished: + # don't add new points if finished + return + + # migrate SceneraioScene function + start_frame_idx = max(self.ptrack.frame_index, self._last_diff_frame_idx) + traj_diff_steps_back = track.frame_index - start_frame_idx # positive value + pred_diff_steps_forward = start_frame_idx - self.ptrack.frame_index # positive value + + if traj_diff_steps_back < 0 or len(track.history) < traj_diff_steps_back: + logger.warning("Track history doesn't reach prediction start. Should not be possible. Skip") + # elif len(ptrack.predictions[0]) < pred_diff_steps_back: + # logger.warning("Prediction does not reach prediction start. Should not be possible. Skip") + else: + trajectory = track.projected_history + + # from start to as far as it gets + trajectory_range = trajectory[-1*traj_diff_steps_back:] + prediction_range = self.ptrack.predictions[0][pred_diff_steps_forward:] # in world coordinate space + line = [] + for i, (p1, p2) in enumerate(zip(trajectory_range, prediction_range)): + diff = (p1[0]-p2[0], p1[1]-p2[1]) + self.score += np.linalg.norm(diff) + + offset_from_start = (pred_diff_steps_forward + i) + if offset_from_start % self.POINT_INTERVAL == 0: + self.line.extend([p1, p2]) + self.points.extend([p1, p2]) + + self._last_diff_frame_idx = track.frame_index + + + # # run each render tick + # def update_drawn_positions(self, dt: DeltaT): + # if isinstance(self.line, AppendableLine): + # if self.finished and self.line.ready: + # # convert when fully drawn + # # print(self, "CONVERT LINE") + # self.line = ProceduralChain.from_appendable_line(self.line) + + # if isinstance(self.line, ProceduralChain): + # self.line.target = self._target_track.projected_history[-1] + + # # if not self.finished or not self.line.ready: + # self.line.update_drawn_positions(dt) + + + + def as_renderable(self) -> RenderableLines: + color = SrgbaColor(0,0,1,1) + + # if not self.finished or not self.line.ready: + return self.line.as_renderable(color) + # return self.line.as_renderable(color) + + +def calculate_loitering_scores(track: ProjectedTrack, min_duration_to_linger, linger_factor, velocity_threshold, window = None): + """ + Calculates a loitering score (0-1) for each track. + + Args: + tracks: A list of tracks, where each track is a list of (frame_id, x, y, width, height). + min_duration_to_linger: Minimum number of frames to start considering a segment as lingering. + linger_factor: Divide number of lingering frames by 'linger_factor' to get a score 0-1 + velocity_threshold: Maximum velocity (meters/frame) to consider as lingering. + + Returns: + A generator providing loitering scores for each frame + """ + + total_frames = len(track.projected_history) + + if total_frames < 2: + return 0.0 # Not enough data + + offset = window * -1 if window is not None else 0 + + x_coords = [t[0] for t in track.projected_history[offset:]] + y_coords = [t[1] for t in track.projected_history[offset:]] + + # Calculate velocities + velocities = np.sqrt(np.diff(x_coords)**2 + np.diff(y_coords)**2) + + # Calculate distances + # distances = np.diff(x_coords) + # distances_y = np.diff(y_coords) + # distances_total = np.sqrt(distances**2 + distances_y**2) + + linger_duration = 0 + linger_frames = 0 + + + for i in range(len(velocities)): + if velocities[i] < velocity_threshold: + linger_duration += 1 + if linger_duration >= min_duration_to_linger: + linger_frames +=1 + else: + # decay if moving faster + linger_duration = max(linger_duration - 3, 0) + linger_frames = max(linger_frames - 3, 0) + + # Calculate loitering score + if total_frames > 0: + loitering_score = min(1, max(0, linger_frames / linger_factor)) + else: + loitering_score = 0.0 + + yield loitering_score + diff --git a/trap/base.py b/trap/base.py index 75a9fbf..abcd0f4 100644 --- a/trap/base.py +++ b/trap/base.py @@ -25,6 +25,7 @@ from bytetracker.basetrack import TrackState as ByteTrackTrackState import pandas as pd from shapely import Point +from trap.anomaly import calculate_loitering_scores from trap.utils import get_bins, inv_lerp, lerp from trajectron.environment import Environment, Node, Scene from urllib.parse import urlparse diff --git a/trap/lines.py b/trap/lines.py index 4c8ab97..d66cce8 100644 --- a/trap/lines.py +++ b/trap/lines.py @@ -456,7 +456,7 @@ class LineAnimator(StaticLine): # print(self, target_line, bool(target_line), target_line is not None) self.target = target_line if target_line is not None else StaticLine() self.ready = len(self.target) == 0 - self.start_t = time.time() + self.start_t = time.perf_counter() self.skip = False def extend(self, coords): @@ -478,15 +478,20 @@ class LineAnimator(StaticLine): def is_ready(self): return (self.ready or self.skip) and self.target.is_ready() + def start(self): self.target.start() - self.start_t = time.time() + self.start_t = time.perf_counter() return True + + def is_running(self): + # when ready, consider not running + return bool(self.start_t) and not self.is_ready() def running_for(self): if self.start_t: - return time.time() - self.start_t + return time.perf_counter() - self.start_t return 0. @@ -619,6 +624,77 @@ class CropLine(LineAnimator): def apply(self, target_line: RenderableLine, dt: DeltaT) -> RenderableLine: target_line.points = target_line.points[-1 * self.max_points:] return target_line + + +class CropAnimationLine(LineAnimator): + """ + Similar to segment, but on points instead of lenght, and as animation + """ + def __init__(self, target_line = None, max_points = 200, assume_fps=12): + super().__init__(target_line) + self.max_points = max_points + self.assume_fps = assume_fps + # self.ready = True # static filter, always ready + + # def set_frame_offset(self, frame_offset: int): + # self.frame_offset = frame_offset + + def apply(self, target_line: RenderableLine, dt: DeltaT) -> RenderableLine: + + dt = self.running_for() + frame_offset = int(dt * self.assume_fps) + + max_points = int(self.max_points(dt)) if callable(self.max_points) else self.max_points + + + head = frame_offset + 1 + tail = max(0, frame_offset + 1 - max_points) + # print(self.running_for(), frame_offset, head,tail) + + + target_line.points = target_line.points[tail:head] + self.ready = len(target_line.points) < 1 + + return target_line + + +class FadedEndsLine(LineAnimator): + """ + Static filter; Fade the tail of the line. Always (not only cropped) + """ + def __init__(self, target_line = None, in_fade_steps: int = 30, out_fade_steps: int = 30): + super().__init__(target_line) + self.ready = True + self.fade_in_steps = in_fade_steps + self.fade_out_steps = out_fade_steps + + + def apply(self, target_line: RenderableLine, dt: DeltaT) -> RenderableLine: + l = len(target_line.points) + points = [] + + # TODO: fractional divide if fade_in and out are not equal + half_points = l // 2 + fade_in = min(self.fade_in_steps, half_points) + fade_out = min(self.fade_out_steps, half_points) + + + for i, point in enumerate(target_line.points): + if i < fade_in: + t = i / self.fade_in_steps + elif i > (l - fade_out): + t = 1 - (i - (l - fade_out)) / self.fade_out_steps + else: + t = 1 + + alpha = t + if alpha >= 0: + alpha = min(1, alpha) + point.color = point.color.as_faded(alpha) + points.append(point) + + + return RenderableLine(points) class FadedTailLine(LineAnimator): """ @@ -984,6 +1060,9 @@ class LineAnimationStack(): def is_ready(self): return self.tail.is_ready() + def is_running(self): + return self.tail.is_running() + class LineAnimationSequenceStep(NamedTuple): line: LineAnimator @@ -1013,7 +1092,7 @@ class LineAnimationSequence(): def start(self): - self.start_at = time.time() + self.start_at = time.perf_counter() self.idx = 0 def as_renderable_line(self, dt: DeltaT): diff --git a/trap/renderable.proto b/trap/renderable.proto new file mode 100644 index 0000000..aec2e3c --- /dev/null +++ b/trap/renderable.proto @@ -0,0 +1,50 @@ +syntax = "proto3"; + +package renderable; + +// Enum for coordinate spaces +enum CoordinateSpace { + UNDEFINED=0; + CAMERA = 1; + UNDISTORTED_CAMERA = 2; + WORLD = 3; + LASER = 4; + RAW_LASER = 8; +} + +// Message for RenderablePosition (Tuple[float, float]) +message RenderablePosition { + float x = 1; + float y = 2; +} + +// Message for SrgbaColor +message SrgbaColor { + float red = 1; + float green = 2; + float blue = 3; + float alpha = 4; +} + +// Message for RenderablePoint +message RenderablePoint { + RenderablePosition position = 1; + SrgbaColor color = 2; +} + +// Message for RenderableLine +message RenderableLine { + repeated RenderablePoint points = 1; +} + +// Message for RenderableLines +message RenderableLines { + repeated RenderableLine lines = 1; + CoordinateSpace space = 2; +} + +// Message to represent RenderableLayers (Dict[int, RenderableLines]) +message RenderableLayers { + map layers = 1; +} + diff --git a/trap/stage.py b/trap/stage.py index 1b001ad..c7a0776 100644 --- a/trap/stage.py +++ b/trap/stage.py @@ -8,7 +8,9 @@ from functools import partial import json import logging from math import inf +import math from pathlib import Path +import random import time import threading from typing import Dict, Generator, List, Optional, Type, TypeVar @@ -19,7 +21,7 @@ import zmq from trap.anomaly import DiffSegment, calc_anomaly, calculate_loitering_scores from trap.base import CameraAction, DataclassJSONEncoder, Frame, HomographyAction, ProjectedTrack, Track from trap.counter import CounterSender -from trap.lines import AppendableLine, AppendableLineAnimator, Coordinate, CropLine, DashedLine, DeltaT, FadeOutJitterLine, FadeOutLine, FadedTailLine, LineAnimationStack, LineAnimator, NoiseLine, RenderableLayers, RenderableLine, RenderableLines, SegmentLine, SimplifyMethod, SrgbaColor, StaticLine, layers_to_message, load_lines_from_svg +from trap.lines import AppendableLine, AppendableLineAnimator, Coordinate, CoordinateSpace, CropAnimationLine, CropLine, DashedLine, DeltaT, FadeOutJitterLine, FadeOutLine, FadedEndsLine, FadedTailLine, LineAnimationStack, LineAnimator, NoiseLine, RenderableLayers, RenderableLine, RenderableLines, SegmentLine, SimplifyMethod, SrgbaColor, StaticLine, layers_to_message, load_lines_from_svg from trap.node import Node from trap.track_history import TrackHistory @@ -98,6 +100,8 @@ class Scenario: self.is_running = False + self.loitering_factor = 0 + logger.info(f"Found {self.track_id}: {self.scene.name}") def start(self): @@ -114,7 +118,7 @@ class Scenario: if self.take_over_at: return - self.take_over_at = time.time() + self.take_over_at = time.perf_counter() def taken_over(self): self.is_running = False @@ -122,7 +126,7 @@ class Scenario: def takenover_for(self): if self.take_over_at: - return time.time() - self.take_over_at + return time.perf_counter() - self.take_over_at return None def takeover_factor(self): @@ -134,7 +138,7 @@ class Scenario: def lost_for(self): if self.scene is ScenarioScene.LOST: - return time.time() - self.state_change_at + return time.perf_counter() - self.state_change_at return None def lost_factor(self): @@ -144,7 +148,7 @@ class Scenario: return l/LOST_FADEOUT def anomaly_factor(self): - return calc_anomaly(self.prediction_diffs, 10) + return calc_anomaly(self.prediction_diffs) def deactivate(self): self.take_over_at = None @@ -161,7 +165,7 @@ class Scenario: logger.info(f"Changing scene for {self.track_id}: {self.scene.name} -> {scene.name}") self.scene = scene - self.state_change_at = time.time() + self.state_change_at = time.perf_counter() def update_state(self): self.check_lost() or self.check_loitering() or self.check_track() @@ -175,7 +179,9 @@ class Scenario: def check_loitering(self): scores = [s for s in calculate_loitering_scores(self.track, LOITERING_DURATION_TO_LINGER, LOITERING_LINGER_FACTOR, LOITERING_VELOCITY_TRESHOLD/TRACK_ASSUMED_FPS, 150)] - if scores[-1] > .99: + + self.loitering_factor = scores[-1] + if self.loitering_factor > .99: self.set_scene(ScenarioScene.LOITERING) return True return False @@ -284,9 +290,12 @@ class DrawnScenario(Scenario): # when rendering tracks from others similar/close to the current one self.others_color = SrgbaColor(1,1,0,1) self.line_others = LineAnimationStack(StaticLine([], self.others_color)) - self.line_others.add(SegmentLine(self.line_others.tail, duration=3, anim_f=partial(SegmentLine.anim_grow, in_and_out=True, max_len=5))) + # self.line_others.add(SegmentLine(self.line_others.tail, duration=3, anim_f=partial(SegmentLine.anim_grow, in_and_out=True, max_len=5))) + self.line_others.add(CropAnimationLine(self.line_others.tail, lambda dt: 10 + math.sin(dt/4) * 70, assume_fps=TRACK_ASSUMED_FPS*2)) # speed up + self.line_others.add(NoiseLine(self.line_others.tail, amplitude=0, t_factor=.3)) # self.line_others.add(DashedLine(self.line_others.tail, t_factor=4, loop_offset=True)) # self.line_others.get(DashedLine).skip = True + self.line_others.add(FadedEndsLine(self.line_others.tail, 30, 30)) self.line_others.add(FadeOutLine(self.line_others.tail)) self.line_others.get(FadeOutLine).set_alpha(0) @@ -340,9 +349,9 @@ class DrawnScenario(Scenario): # special case: LOITERING - if self.scene is ScenarioScene.LOITERING or self.state_change_at: + if self.scene is ScenarioScene.LOITERING: # or self.state_change_at: # logger.info('loitering') - transition = min(1, (time.time() - self.state_change_at)/1.4) + transition = min(1, (time.perf_counter() - self.state_change_at)/1.4) # TODO: transition fade, using to_alpha(), so it can fade back in again: @@ -362,26 +371,38 @@ class DrawnScenario(Scenario): self.tracks_to_self_fetched_at = time.perf_counter() # fetch lines nearby - track_ids = stage.history.get_nearest_tracks(current_position, 15) + track_ids = stage.history.get_nearest_tracks(current_position, 30) self.track_ids_to_self = iter(track_ids) self.tracks_to_self = stage.history.ids_as_trajectory(track_ids) - print(time.perf_counter() - t, "fetch delya") + print(time.perf_counter() - t, "fetch delay") - if self.tracks_to_self and self.line_others.is_ready(): - current_history_id = next(self.track_ids_to_self) - current_history = next(self.tracks_to_self) + # if self.tracks_to_self and not len(self.line_others.root.points): + if self.tracks_to_self and not self.line_others.is_running(): + try: + current_history = next(self.tracks_to_self) + current_history_id = next(self.track_ids_to_self) + self.line_others.get(CropAnimationLine).assume_fps += TRACK_ASSUMED_FPS*1.5 # faster each time + self.line_others.get(NoiseLine).amplitude = .05 - logger.info(f"play history item: {current_history_id}") - self.line_others.get(FadeOutLine).set_alpha(1) - - self.line_others.root.points = current_history - # print(self.line_others.root.points) - self.line_others.start() + logger.info(f"play history item: {current_history_id}") + self.line_others.get(FadeOutLine).set_alpha(1) + + self.line_others.root.points = current_history + # print(self.line_others.root.points) + self.line_others.start() + except StopIteration as e: + pass + # logger.info("Exhausted similar tracks?") + + else: + # reset loitering values + self.line_others.get(CropAnimationLine).assume_fps = TRACK_ASSUMED_FPS*2 + self.line_others.get(NoiseLine).amplitude = 0 # special case: PLAY - elif self.scene is ScenarioScene.PLAY: + if self.scene is ScenarioScene.PLAY: pass # if self.scene is ScenarioScene.CORRECTED_PREDICTION: # self.line_prediction.get(DashedLine).skip = False @@ -391,8 +412,10 @@ class DrawnScenario(Scenario): def to_renderable_lines(self, dt: DeltaT) -> RenderableLines: # each scene is handled differently: + t1 = time.perf_counter() # 1) history, fade out when lost + # self.line_history.get(StaticLine).color = SrgbaColor(1, 0, 1-self.anomaly_factor(), 1) self.line_history.get(FadeOutJitterLine).set_alpha(1-self.lost_factor()) self.line_prediction.get(FadeOutLine).set_alpha(1-self.lost_factor()) self.line_history.get(NoiseLine).amplitude = self.lost_factor() @@ -401,23 +424,60 @@ class DrawnScenario(Scenario): track_age_in_frames = self.track_age() * TRACK_ASSUMED_FPS self.line_history.get(FadedTailLine).set_frame_offset(track_age_in_frames) - + t2 = time.perf_counter() # 2) also fade-out when moving into loitering mode. # when fading out is done, start drawing historical data history_line = self.line_history.as_renderable_line(dt) + t3 = time.perf_counter() prediction_line = self.line_prediction.as_renderable_line(dt) + t4 = time.perf_counter() others_line = self.line_others.as_renderable_line(dt) + t5 = time.perf_counter() # print(history_line) # print(self.track_id, len(self.line_history.points), len(history_line)) + timings = (t5-t4, t4-t3, t3-t2, t2-t1) + return RenderableLines([ history_line, prediction_line, others_line - ]) + ]), timings + + +class DatasetDrawer(): + def __init__(self, stage: Stage): + self.stage = stage + + line_color = SrgbaColor(0,1,1,1) + self.track_line = LineAnimationStack(StaticLine([], line_color)) + self.track_line.add(CropAnimationLine(self.track_line.tail, 100, assume_fps=TRACK_ASSUMED_FPS*30)) # speed up + # self.track_line.add(DashedLine(self.track_line.tail, t_factor=4, loop_offset=True)) + # self.track_line.get(DashedLine).skip = True + self.track_line.add(FadedEndsLine(self.track_line.tail, 10, 10)) + # self.track_line.add(FadeOutLine(self.track_line.tail)) + # self.track_line.get(FadeOutLine).set_alpha(1) + + def to_renderable_lines(self, dt: DeltaT): + lines = RenderableLines([], CoordinateSpace.WORLD) + if not self.track_line.is_running(): + track_id = random.choice(list(self.stage.history.state.tracks.keys())) + # print('track_id', track_id) + positions = self.stage.history.state.track_histories[track_id] + self.track_line.root.points = positions + self.track_line.start() + + lines.lines.append( + self.track_line.as_renderable_line(dt) + ) + # print(lines) + + return lines + + class Stage(Node): @@ -441,6 +501,8 @@ class Stage(Node): self.debug_lines = RenderableLines(load_lines_from_svg(self.config.debug_map, 100, debug_color)) self.history = TrackHistory(self.config.tracker_output_dir, self.config.camera, self.config.cache_path) + + self.auxilary = DatasetDrawer(self) @@ -536,10 +598,15 @@ class Stage(Node): # TODO: sometimes very slow! t1 = time.perf_counter() + timings = [] for scenario in self.active_scenarios: - lines.append_lines(scenario.to_renderable_lines(dt)) + scenario_lines, timing = scenario.to_renderable_lines(dt) + lines.append_lines(scenario_lines) + timings.append(timing) t2 = time.perf_counter() + training_lines = self.auxilary.to_renderable_lines(dt) + t2b = time.perf_counter() rl = lines.as_simplified(SimplifyMethod.RDP, .003) # or segmentise (see shapely) self.counter.set("stage.lines", len(lines.lines)) self.counter.set("stage.points_orig", lines.point_count()) @@ -550,6 +617,7 @@ class Stage(Node): layers: RenderableLayers = { 1: lines, 2: self.debug_lines, + 3: training_lines, } t4 = time.perf_counter() @@ -564,12 +632,13 @@ class Stage(Node): t6 = time.perf_counter() - t = (t2-t1, t3-t2, t4-t3, t5-t4, t6-t5) + t = (t2-t1, t3-t2b, t2b-t2, t4-t3, t5-t4, t6-t5) if sum(t) > .1: print(t) print(len(lines.lines)) print(lines.point_count()) print(len(msg)) + print('scenario timings:', timings) # print(msg) # exit() diff --git a/trap/track_history.py b/trap/track_history.py new file mode 100644 index 0000000..c3c4efb --- /dev/null +++ b/trap/track_history.py @@ -0,0 +1,183 @@ +from dataclasses import dataclass +import logging +from pathlib import Path +import pickle +from threading import Lock +import time +from typing import Dict, Iterable, List, Optional, Set + +import numpy as np +from trap.base import Camera, Track +from trap.lines import Coordinate +from trap.tracker import FinalDisplacementFilter, Smoother, TrackReader + +from scipy.spatial import KDTree + +logger = logging.getLogger('history') + +@dataclass +class TrackHistoryState(): + """ + The lock of TrackHistory is not pickle-able so separate it into a separate state + """ + tracks: List[Track] + track_histories: Dict[str, np.ndarray] + indexed_track_ids: List[str] + tree: KDTree + + + +class TrackHistory(): + def __init__(self, path: Path, camera: Camera, cache_path: Optional[Path]): + self.path = path + self.camera = camera + self.cache_path = cache_path + self.lock = Lock() + self.load_from_cache() or self.reload() + + + def load_from_cache(self): + if self.cache_path.exists(): + logger.debug("Load history state from cache") + with self.cache_path.open('rb') as fp: + try: + state = pickle.load(fp) + if not isinstance(state, TrackHistoryState): + raise RuntimeError("Pickled data is not a trackhistorystate") + self.state = state + return True + except Exception as e: + logger.warning(f"Cannot read cache {self.cache_path}: {e}") + + return False + + def build_tree(self): + reader = TrackReader(self.path, self.camera.fps) + logger.debug(f'loaded {len(reader)} tracks') + + track_filter = FinalDisplacementFilter(2) + tracks = track_filter.apply(reader, self.camera) + logger.debug(f'after filtering left with {len(tracks)} tracks') + + + tracks: List[Track] = [t.get_with_interpolated_history() for t in tracks] + logger.debug(f'interpolated {len(tracks)} tracks') + + # use convolution here, because precision does not matter and it is _way_ faster + smoother = Smoother(convolution=True) + tracks = [smoother.smooth_track(t) for t in tracks] + logger.debug(f'smoothed') + + tracks = {track.track_id: track for track in tracks} + + + track_histories = {t.track_id: t.get_projected_history(camera=self.camera) for t in tracks.values()} + downsampled_histories = {t_id: self.downsample_history(h) for t_id, h in track_histories.items()} + logger.debug(f'projected to world space') + + + # Sample data (coordinates and metadata) + # coordinates = [(1, 2, 'Point A'), (3, 4, 'Point B'), (5, 6, 'Point C'), (7, 8, 'Point D')] + all_points = [] + indexed_track_ids: List[str] = [] + for track_id, history in downsampled_histories.items(): + all_points.extend([ + [point[0], point[1]] for point in history + ]) + indexed_track_ids.extend([track_id] * len(history)) + + # self.flat_idx = self.flat_histories[:,2] + + # Create the KD-Tree + tree = KDTree(all_points) + + logger.debug('built tree') + return TrackHistoryState( + tracks, track_histories, indexed_track_ids, tree + ) + + def reload(self): + state = self.build_tree() + + # aquire lock as brief as possible + with self.lock: + self.state = state + + + if self.cache_path: + with self.cache_path.open('wb') as fp: + logger.debug("Writing history to cache") + pickle.dump(self.state, fp) + + + + def get_nearest_tracks(self, point: Coordinate, k:int, max_r: Optional[float] = np.inf): + with self.lock: + distances, indexes = self.state.tree.query(point, k, distance_upper_bound=max_r) + # filter out when there's no + indexes = indexes[distances != np.inf] + track_ids: Set[str] = {self.state.indexed_track_ids[idx] for idx in indexes} + + # nearby_indexes = self.tree.query_ball_point(point, r) + # track_ids = set([self.flat_idx[idx] for idx in nearby_indexes]) + + return track_ids + + def ids_as_trajectory(self, track_ids: Iterable[str]): + for track_id in track_ids: + yield self.state.tracks[track_id].get_projected_history(camera=self.camera) + + + + @classmethod + def downsample_history(cls, history, cell_size=.3): + + + if not len(history): + return [] + + positions = np.unique(np.round(history / cell_size), axis=0) * cell_size + + return positions + + + + + +if __name__ == "__main__": + path = Path("EXPERIMENTS/raw/hof3/") + logging.basicConfig(level=logging.DEBUG) + + calibration_path = Path("../DATASETS/hof3/calibration.json") + homography_path = Path("../DATASETS/hof3/homography.json") + camera = Camera.from_paths(calibration_path, homography_path, 12) + # device = device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + + + s = time.time() + history = TrackHistory(path, camera, Path("/tmp/historystate_hof3.pcl")) + dt = time.time() - s + print(f'loaded {len(history.state.tracks)} tracks in {dt}s') + + + track = list(history.state.tracks.values())[25] + trajectory_crop = TrackHistory.downsample_history(history.state.track_histories[track.track_id]) + trajectory_org = track.get_projected_history(camera=camera) + target_point = trajectory_org[len(trajectory_org)//2+90] + + import matplotlib.pyplot as plt # Visualization + + track_set = history.get_nearest_tracks(target_point, 10, max_r=np.inf) + + + + plt.gca().set_aspect('equal') + plt.scatter(trajectory_crop[:,0], trajectory_crop[:,1], c='orange') + plt.plot(trajectory_org[:,0], trajectory_org[:,1], c='blue', alpha=1) + plt.scatter(target_point[0], target_point[1], c='red', alpha=1) + for track_id in track_set: + closeby = history.state.tracks[track_id].get_projected_history(camera=camera) + plt.plot(closeby[:,0], closeby[:,1], c='green', alpha=.1) + + plt.show() \ No newline at end of file