diff --git a/trap/stage.py b/trap/stage.py index ee96c99..df30427 100644 --- a/trap/stage.py +++ b/trap/stage.py @@ -16,7 +16,7 @@ import zmq from sgan.sgan import data from trap import shapes -from trap.base import DataclassJSONEncoder, Frame, Track +from trap.base import Camera, DataclassJSONEncoder, DistortedCamera, Frame, Track from trap.counter import CounterSender from trap.laser_renderer import circle_points, rotateMatrix from trap.lines import RenderableLine, RenderableLines, RenderablePoint, SrgbaColor, circle_arc @@ -25,6 +25,8 @@ from trap.timer import Timer from trap.utils import exponentialDecay, relativePointToPolar, relativePolarToPoint +logger = logging.getLogger('trap.stage') + class ScenarioScene(Enum): DETECTED = 1 FIRST_PREDICTION = 2 @@ -34,7 +36,7 @@ class ScenarioScene(Enum): LOST = -1 LOST_FADEOUT = 3 -PREDICTION_INTERVAL: float|None = 15 # frames +PREDICTION_INTERVAL: float|None = 20 # frames PREDICTION_FADE_IN: float = 3 PREDICTION_FADE_SLOPE: float = -10 PREDICTION_FADE_AFTER_DURATION: float = 10 # seconds @@ -247,6 +249,7 @@ class DrawnScenario(TrackScenario): # 3. predictions self.drawn_predictions = [] + self.drawn_diffs = [] for a, (ptrack, next_ptrack) in enumerate(zip(self._predictions, [*self._predictions[1:], None])): prediction = ptrack.predictions[0] # only use one prediction per timestep/frame/track @@ -254,10 +257,29 @@ class DrawnScenario(TrackScenario): # not the last one, cut off next_ptrack: Track = self._predictions[a+1] end_step = next_ptrack.frame_index - ptrack.frame_index + + # diff + diff_steps_back = ptrack.frame_index - self._track.frame_index + if len(self.drawn_positions) < -1*diff_steps_back: + logger.warning("Track history doesn't reach prediction start. Should not be possible. Skip") + pass + else: + # trajectory_range = self.camera.[d.get_foot_coords() for d in trajectory_det_range] # in frame coordinate space + trajectory_range = self.drawn_positions[diff_steps_back:diff_steps_back+end_step] + prediction_range = ptrack.predictions[0][:end_step] # in world coordinate space + line = [] + for p1, p2 in zip(trajectory_range[::4], prediction_range[::4]): + line.extend([ + p1, p2 + ]) + if len(line): + self.drawn_diffs.append(line) else: end_step = None # not last item; show all self.drawn_predictions.append(ptrack.predictions[0][:end_step]) + + @@ -340,6 +362,11 @@ class DrawnScenario(TrackScenario): # points = [RenderablePoint(pos, pos_color) for pos, pos_color in zip(drawn_prediction[PREDICTION_OFFSET:], colors[PREDICTION_OFFSET:])] points = [RenderablePoint(pos, pos_color) for pos, pos_color in zip(drawn_prediction, colors)] lines.append(RenderableLine(points)) + for drawn_diff in self.drawn_diffs: + color = SrgbaColor(0.,1,1.,1.-self.lost_factor()) + colors = [color.as_faded(1) for a2 in range(len(drawn_diff))] + points = [RenderablePoint(pos, pos_color) for pos, pos_color in zip(drawn_diff, colors)] + lines.append(RenderableLine(points)) # # print(self.current_state) # if self.current_state is self.first_prediction or self.current_state is self.corrected_prediction: @@ -460,6 +487,7 @@ class Stage(Node): self.stage_sock = self.pub(self.config.zmq_stage_addr) self.counter = CounterSender() + self.camera: Optional[DistortedCamera] = None def run(self):