trap/config.py

@@ -255,9 +255,6 @@ render_parser.add_argument("--render-file",
 render_parser.add_argument("--render-window",
                     help="Render a previewing to a window",
                     action='store_true')
-render_parser.add_argument("--full-screen",
-                    help="Set Window full screen",
-                    action='store_true')
 
 render_parser.add_argument("--render-url",
                     help="""Stream renderer on given URL. Two easy approaches:

trap/renderer.py

@@ -18,7 +18,6 @@ import tempfile
 from pathlib import Path
 import shutil
 import math
-from collections import deque
 
 from pyglet import shapes
 from PIL import Image
@@ -59,21 +58,15 @@ def relativePolarToPoint(origin, r, angle) -> tuple[float, float]:
 
 class DrawnTrack:
     def __init__(self, track_id, track: Track, renderer: Renderer, H):
-        # self.created_at = time.time()
-        self.update_at = self.created_at = time.time()
         self.track_id = track_id
         self.renderer = renderer
+        self.set_track(track, H)
         self.drawn_positions = []
         self.drawn_predictions = []
-        self.predictions: deque[DrawnPrediction] = deque(maxlen=20) # TODO; make configurable
         self.shapes: list[pyglet.shapes.Line] = []
-
+        self.pred_shapes: list[list[pyglet.shapes.Line]] = []
-        self.set_track(track, H)
-        self.set_prediction(track)
 
     def set_track(self, track: Track, H):
-        self.update_at = time.time()
-
         self.track = track
         self.H = H
         self.coords = [d.get_foot_coords() for d in track.history]
@@ -81,17 +74,11 @@ class DrawnTrack:
         # perhaps only do in constructor:
         self.inv_H = np.linalg.pinv(self.H)
 
-    def set_prediction(self, track: Track):
-        # TODO: turn into add_prediction
         pred_coords = []
-        if not track.predictions:
-            return
-        
         for pred_i, pred in enumerate(track.predictions):
             pred_coords.append(cv2.perspectiveTransform(np.array([pred]), self.inv_H)[0].tolist())
         
-        # self.pred_coords = pred_coords
+        self.pred_coords = pred_coords
-        self.predictions.append(DrawnPrediction(self, pred_coords))
             # color = (128,0,128) if pred_i else (128,
         
         
@@ -107,26 +94,23 @@ class DrawnTrack:
         if len(self.coords) > len(self.drawn_positions):
             self.drawn_positions.extend(self.coords[len(self.drawn_positions):])
 
-        # Superseded by individual drawnprediction elements
+        for a, drawn_prediction in enumerate(self.drawn_predictions):
-        # for a, drawn_prediction in enumerate(self.drawn_predictions):
+            for i, pos in enumerate(drawn_prediction):
-        #     for i, pos in enumerate(drawn_prediction):
+                # TODO: this should be done in polar space starting from origin (i.e. self.drawn_posision[-1])
-        #         # TODO: this should be done in polar space starting from origin (i.e. self.drawn_posision[-1])
+                decay = max(3, (18/i) if i else 10) # points further away move with more delay
-        #         decay = max(3, (18/i) if i else 10) # points further away move with more delay
+                decay = 6
-        #         decay = 6
+                origin = self.drawn_positions[-1]
-        #         origin = self.drawn_positions[-1]
+                drawn_r, drawn_angle = relativePointToPolar(  origin, drawn_prediction[i])
-        #         drawn_r, drawn_angle = relativePointToPolar(  origin, drawn_prediction[i])
+                pred_r, pred_angle = relativePointToPolar(origin, self.pred_coords[a][i])
-        #         pred_r, pred_angle = relativePointToPolar(origin, self.pred_coords[a][i])
+                r = exponentialDecay(drawn_r, pred_r, decay, dt)
-        #         r = exponentialDecay(drawn_r, pred_r, decay, dt)
+                angle = exponentialDecay(drawn_angle, pred_angle, decay, dt)
-        #         angle = exponentialDecay(drawn_angle, pred_angle, decay, dt)
+                x, y = relativePolarToPoint(origin, r, angle)
-        #         x, y = relativePolarToPoint(origin, r, angle)
+                self.drawn_predictions[a][i] = int(x), int(y)
-        #         self.drawn_predictions[a][i] = int(x), int(y)
+                # self.drawn_predictions[i][0] = int(exponentialDecay(self.drawn_predictions[i][0], self.pred_coords[i][0], decay, dt))
-        #         # self.drawn_predictions[i][0] = int(exponentialDecay(self.drawn_predictions[i][0], self.pred_coords[i][0], decay, dt))
+                # self.drawn_predictions[i][1] = int(exponentialDecay(self.drawn_predictions[i][1], self.pred_coords[i][1], decay, dt))
-        #         # self.drawn_predictions[i][1] = int(exponentialDecay(self.drawn_predictions[i][1], self.pred_coords[i][1], decay, dt))
-
-        #     if len(self.pred_coords) > len(self.drawn_predictions):
-        #         self.drawn_predictions.extend(self.pred_coords[len(self.drawn_predictions):])
-
 
+        if len(self.pred_coords) > len(self.drawn_predictions):
+            self.drawn_predictions.extend(self.pred_coords[len(self.drawn_predictions):])
         # for a, drawn_prediction in self.drawn_predictions:
         #     if len(self.pred_coords) > len(self.drawn_predictions):
         #         self.drawn_predictions.extend(self.pred_coords[len(self.drawn_predictions):])
@@ -204,55 +188,6 @@ class DrawnTrack:
                         target_opacity = (1 - ((ci - half) / half)) * 180
                     line.opacity = int(exponentialDecay(line.opacity, target_opacity, decay, dt))
         
-class DrawnPrediction:
-    def __init__(self, drawn_track: DrawnTrack, coords: list[list] = []):
-        self.created_at = time.time()
-        # self.renderer = renderer
-        self.drawn_track = drawn_track
-        self.coords = coords
-        self.color = colorset[self.drawn_track.track_id % len(colorset)]
-        self.pred_shapes: list[list[pyglet.shapes.Line]] = []
-
-        # coords is a list of predictions
-        for a, coords in enumerate(self.coords):
-            prediction_shapes = []
-            for ci in range(0, len(coords)):
-                if ci == 0:
-                    x, y = [int(p) for p in self.drawn_track.coords[-1]]
-                else:
-                    x, y = [int(p) for p in coords[ci-1]]
-
-                x2, y2 = [int(p) for p in coords[ci]]
-
-                # flip in window:
-                y, y2 = self.drawn_track.renderer.window.height - y, self.drawn_track.renderer.window.height - y2
-                
-                line = self.drawn_track.renderer.gradientLine(x, y, x2, y2, 3, self.color, self.color, batch=self.drawn_track.renderer.batch_anim)
-                line.opacity = 5
-                prediction_shapes.append(line)
-            self.pred_shapes.append(prediction_shapes)
-
-
-    def update_opacities(self, dt: float):
-        """
-        Update the opacties of the drawn line, by only using the dt provided by the renderer
-        Done using exponential decal, with a different decay value per item
-        """
-        for a, coords in enumerate(self.coords):
-            for ci in range(0, len(coords)):
-                line = self.pred_shapes[a][ci-1]
-                # Positions of prediction no longer update
-                # line.x, line.y = x, y
-                # line.x2, line.y2 = x2, y2
-                # line.color = color
-                decay = (16/ci) if ci else 16
-                half = len(coords) / 2
-                if ci < half:
-                    target_opacity = 180
-                else:
-                    target_opacity = (1 - ((ci - half) / half)) * 180
-                line.opacity = int(exponentialDecay(line.opacity, target_opacity, decay, dt))
-        
 
 class FrameWriter:
     """
@@ -304,11 +239,6 @@ class Renderer:
         self.prediction_sock.setsockopt(zmq.SUBSCRIBE, b'')
         self.prediction_sock.connect(config.zmq_prediction_addr if not self.config.bypass_prediction else config.zmq_trajectory_addr)
         
-        self.tracker_sock = context.socket(zmq.SUB)
-        self.tracker_sock.setsockopt(zmq.CONFLATE, 1) # only keep latest frame. NB. make sure this comes BEFORE connect, otherwise it's ignored!!
-        self.tracker_sock.setsockopt(zmq.SUBSCRIBE, b'')
-        self.tracker_sock.connect(config.zmq_trajectory_addr)
-        
         self.frame_sock = context.socket(zmq.SUB)
         self.frame_sock.setsockopt(zmq.CONFLATE, 1) # only keep latest frame. NB. make sure this comes BEFORE connect, otherwise it's ignored!!
         self.frame_sock.setsockopt(zmq.SUBSCRIBE, b'')
@@ -335,7 +265,7 @@ class Renderer:
         
         config = pyglet.gl.Config(sample_buffers=1, samples=4)
         # , fullscreen=self.config.render_window
-        self.window = pyglet.window.Window(width=self.frame_size[0], height=self.frame_size[1], config=config, fullscreen=self.config.full_screen)
+        self.window = pyglet.window.Window(width=self.frame_size[0], height=self.frame_size[1], config=config)
         self.window.set_handler('on_draw', self.on_draw)
         self.window.set_handler('on_refresh', self.on_refresh)
         self.window.set_handler('on_close', self.on_close)
@@ -352,7 +282,6 @@ class Renderer:
 
         self.first_time: float|None = None
         self.frame: Frame|None= None
-        self.tracker_frame: Frame|None = None
         self.prediction_frame: Frame|None = None
 
 
@@ -412,9 +341,7 @@ class Renderer:
 
     def init_labels(self):
         base_color = (255,)*4
-        color_predictor = (255,255,0, 255)
+        info_color = (255,255,0, 255)
-        color_info = (255,0, 255, 255)
-        color_tracker = (0,255, 255, 255)
 
         options = []
         for option in ['prediction_horizon','num_samples','full_dist','gmm_mode','z_mode', 'model_dir']:
@@ -423,15 +350,13 @@ class Renderer:
         self.labels = {
             'waiting': pyglet.text.Label("Waiting for prediction"),
             'frame_idx': pyglet.text.Label("", x=20, y=self.window.height - 17, color=base_color, batch=self.batch_overlay),
-            'tracker_idx': pyglet.text.Label("", x=90, y=self.window.height - 17, color=color_tracker, batch=self.batch_overlay),
+            'frame_time': pyglet.text.Label("t", x=120, y=self.window.height - 17, color=base_color, batch=self.batch_overlay),
-            'pred_idx': pyglet.text.Label("", x=110, y=self.window.height - 17, color=color_predictor, batch=self.batch_overlay),
+            'pred_idx': pyglet.text.Label("", x=90, y=self.window.height - 17, color=info_color, batch=self.batch_overlay),
-            'frame_time': pyglet.text.Label("t", x=140, y=self.window.height - 17, color=base_color, batch=self.batch_overlay),
+            'pred_time': pyglet.text.Label("", x=200, y=self.window.height - 17, color=info_color, batch=self.batch_overlay),
-            'frame_latency': pyglet.text.Label("", x=235, y=self.window.height - 17, color=color_info, batch=self.batch_overlay),
+            'track_len': pyglet.text.Label("", x=500, y=self.window.height - 17, color=base_color, batch=self.batch_overlay),
-            'tracker_time': pyglet.text.Label("", x=300, y=self.window.height - 17, color=color_tracker, batch=self.batch_overlay),
-            'pred_time': pyglet.text.Label("", x=360, y=self.window.height - 17, color=color_predictor, batch=self.batch_overlay),
-            'track_len': pyglet.text.Label("", x=800, y=self.window.height - 17, color=color_tracker, batch=self.batch_overlay),
             'options1': pyglet.text.Label(options.pop(-1), x=20, y=30, color=base_color, batch=self.batch_overlay),
             'options2': pyglet.text.Label(" | ".join(options), x=20, y=10, color=base_color, batch=self.batch_overlay),
+            
         }
     
     def refresh_labels(self, dt: float):
@@ -439,18 +364,12 @@ class Renderer:
         
         if self.frame:
             self.labels['frame_idx'].text = f"{self.frame.index:06d}"
-            self.labels['frame_time'].text = f"{self.frame.time - self.first_time: >10.2f}s"
+            self.labels['frame_time'].text = f"{self.frame.time - self.first_time:.3f}s"
-            self.labels['frame_latency'].text = f"{self.frame.time - time.time():.2f}s"
-        
-        if self.tracker_frame:
-            self.labels['tracker_idx'].text = f"{self.tracker_frame.index - self.frame.index}"
-            self.labels['tracker_time'].text = f"{self.tracker_frame.time - time.time():.3f}s"
-            self.labels['track_len'].text = f"{len(self.tracker_frame.tracks)} tracks"
         
         if self.prediction_frame:
             self.labels['pred_idx'].text = f"{self.prediction_frame.index - self.frame.index}"
             self.labels['pred_time'].text = f"{self.prediction_frame.time - time.time():.3f}s"
-            # self.labels['track_len'].text = f"{len(self.prediction_frame.tracks)} tracks"
+            self.labels['track_len'].text = f"{len(self.prediction_frame.tracks)} tracks"
 
 
         # cv2.putText(img, f"{frame.index:06d}", (20,17), cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
@@ -481,9 +400,8 @@ class Renderer:
             if not self.first_time:
                 self.first_time = self.frame.time
             img = cv2.GaussianBlur(self.frame.img, (15, 15), 0)
-            img = cv2.cvtColor(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2RGB)
+            img = cv2.flip(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), 0)
-            channels = 3 # unfortunately, pyglet seems to draw single channel as Red only
+            img = pyglet.image.ImageData(self.frame_size[0], self.frame_size[1], 'RGB', img.tobytes())
-            img = pyglet.image.ImageData(self.frame_size[0], self.frame_size[1], 'RGB', img.tobytes(), pitch=self.frame_size[0] * -1 * channels)
             # don't draw in batch, so that it is the background
             self.video_sprite = pyglet.sprite.Sprite(img=img, batch=self.batch_bg)
             self.video_sprite.opacity = 100
@@ -493,49 +411,26 @@ class Renderer:
             pass
         try:
             self.prediction_frame: Frame = self.prediction_sock.recv_pyobj(zmq.NOBLOCK)
-            self.update_predictions()
-        except zmq.ZMQError as e:
-            pass
-        try:
-            self.tracker_frame: Frame = self.tracker_sock.recv_pyobj(zmq.NOBLOCK)
             self.update_tracks()
         except zmq.ZMQError as e:
             pass
 
-
     def update_tracks(self):
         """Updates the track objects and shapes. Called after setting `prediction_frame`
         """
         
-        # clean up
-        # for track_id in list(self.drawn_tracks.keys()):
-        #     if track_id not in self.prediction_frame.tracks.keys():
-        #         # TODO fade out
-        #         del self.drawn_tracks[track_id]
-        
-        if self.tracker_frame:
-            for track_id, track in self.tracker_frame.tracks.items():
-                if track_id not in self.drawn_tracks:
-                    self.drawn_tracks[track_id] = DrawnTrack(track_id, track, self, self.tracker_frame.H)
-                else:
-                    self.drawn_tracks[track_id].set_track(track, self.tracker_frame.H)
-        
         # clean up
         for track_id in list(self.drawn_tracks.keys()):
-            # TODO make delay configurable
+            if track_id not in self.prediction_frame.tracks.keys():
-            if self.drawn_tracks[track_id].update_at < time.time() - 5:
                 # TODO fade out
                 del self.drawn_tracks[track_id]
         
-    def update_predictions(self):
-        if self.prediction_frame:
-            for track_id, track in self.prediction_frame.tracks.items():
-                if track_id not in self.drawn_tracks:
-                    self.drawn_tracks[track_id] = DrawnTrack(track_id, track, self, self.prediction_frame.H)
-                    logger.warning("Prediction for uninitialised frame. This should not happen? (maybe huge delay in prediction?)")
-                else:
-                    self.drawn_tracks[track_id].set_prediction(track)
 
+        for track_id, track in self.prediction_frame.tracks.items():
+            if track_id not in self.drawn_tracks:
+                self.drawn_tracks[track_id] = DrawnTrack(track_id, track, self, self.prediction_frame.H)
+            else:
+                self.drawn_tracks[track_id].set_track(track, self.prediction_frame.H)
             
     
     def on_key_press(self, symbol, modifiers):
@@ -575,10 +470,9 @@ class Renderer:
             for shape in track.shapes:
                 shape.draw() # for some reason the batches don't work
         for track in self.drawn_tracks.values():
-            for prediction in track.predictions:
+            for shapes in track.pred_shapes:
-                for shapes in prediction.pred_shapes:
+                for shape in shapes:
-                    for shape in shapes:
+                    shape.draw()
-                        shape.draw()
         # self.batch_anim.draw()
         self.batch_overlay.draw()
         
@@ -654,7 +548,6 @@ class Renderer:
     def run(self):
         frame = None
         prediction_frame = None
-        tracker_frame = None
 
         i=0
         first_time = None
@@ -731,7 +624,7 @@ colorset = [(0, 0, 0),
  (255, 255, 0)
  ]
 
-# Deprecated
+
 def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, config: Namespace) -> np.array:
     # TODO: replace opencv with QPainter to support alpha? https://doc.qt.io/qtforpython-5/PySide2/QtGui/QPainter.html#PySide2.QtGui.PySide2.QtGui.QPainter.drawImage
     # or https://github.com/pygobject/pycairo?tab=readme-ov-file
@@ -740,12 +633,12 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, con
     # or https://api.arcade.academy/en/stable/index.html (supports gradient color in line -- "Arcade is built on top of Pyglet and OpenGL.")
     frame.img
 
-    # # Fill image with red color(set each pixel to red)
+    overlay = np.zeros(frame.img.shape, np.uint8)
-    # overlay = np.zeros(frame.img.shape, np.uint8)
+    # Fill image with red color(set each pixel to red)
-    # overlay[:] = (130, 0, 75)
+    overlay[:] = (130, 0, 75)
 
-    # img = cv2.addWeighted(frame.img, .4, overlay, .6, 0)
+    img = cv2.addWeighted(frame.img, .4, overlay, .6, 0)
-    img = frame.img.copy()
+    # img = frame.img.copy()
 
     # all not working:
     # if i == 1: