Predictor options now configurable and rendered

trap/config.py

@@ -152,6 +152,31 @@ inference_parser.add_argument('--predict_training_data',
                     help='Ignore tracker and predict data from the training dataset',
                     action='store_true')
 
+inference_parser.add_argument("--smooth-predictions",
+                    help="Smooth the predicted tracks",
+                    action='store_true')
+
+inference_parser.add_argument('--prediction-horizon',
+                    help='Trajectron.incremental_forward parameter',
+                    type=int,
+                    default=30)
+inference_parser.add_argument('--num-samples',
+                    help='Trajectron.incremental_forward parameter',
+                    type=int,
+                    default=5)
+inference_parser.add_argument("--full-dist",
+                    help="Trajectron.incremental_forward parameter",
+                    type=bool,
+                    default=False)
+inference_parser.add_argument("--gmm-mode",
+                    help="Trajectron.incremental_forward parameter",
+                    type=bool,
+                    default=True)
+inference_parser.add_argument("--z-mode",
+                    help="Trajectron.incremental_forward parameter",
+                    type=bool,
+                    default=False)
+
 
 # Internal connections.
 
@@ -192,6 +217,10 @@ frame_emitter_parser.add_argument("--video-src",
                     type=Path,
                     nargs='+',
                     default=lambda: list(Path('../DATASETS/VIRAT_subset_0102x/').glob('*.mp4')))
+frame_emitter_parser.add_argument("--video-offset",
+                    help="Start playback from given frame. Note that when src is an array, this applies to all videos individually.",
+                    default=None,
+                    type=int)
 #TODO: camera as source
 
 frame_emitter_parser.add_argument("--video-loop",
@@ -214,6 +243,9 @@ tracker_parser.add_argument("--detector",
                     help="Specify the detector to use",
                     type=str,
                     choices=DETECTORS)
+tracker_parser.add_argument("--smooth-tracks",
+                    help="Smooth the tracker tracks before sending them to the predictor",
+                    action='store_true')
 
 # Renderer
 

trap/frame_emitter.py

@@ -42,6 +42,7 @@ class Detection:
     h: int # height - image space
     conf: float # object detector probablity
     state: DetectionState
+    frame_nr: int
 
     def get_foot_coords(self):
         return [self.l + 0.5 * self.w, self.t+self.h]
@@ -149,6 +150,12 @@ class FrameEmitter:
             target_frame_duration = 1./fps
             logger.info(f"Emit frames at {fps} fps")
 
+            if self.config.video_offset:
+                logger.info(f"Start at frame {self.config.video_offset}")
+                video.set(cv2.CAP_PROP_POS_FRAMES, self.config.video_offset)
+                i = self.config.video_offset
+
+
             if '-' in video_path.stem:
                 path_stem = video_path.stem[:video_path.stem.rfind('-')]
             else:

trap/prediction_server.py

@@ -27,7 +27,7 @@ import matplotlib.pyplot as plt
 import zmq
 
 from trap.frame_emitter import Frame
-from trap.tracker import Track
+from trap.tracker import Track, Smoother
 
 logger = logging.getLogger("trap.prediction")
 
@@ -121,6 +121,9 @@ class PredictionServer:
         if self.config.eval_device == 'cpu':
             logger.warning("Running on CPU. Specifying --eval_device cuda:0 should dramatically speed up prediction")
 
+        if self.config.smooth_predictions:
+            self.smoother = Smoother(window_len=4)
+        
         context = zmq.Context()
         self.trajectory_socket: zmq.Socket = context.socket(zmq.SUB)
         self.trajectory_socket.setsockopt(zmq.SUBSCRIBE, b'')
@@ -188,7 +191,7 @@ class PredictionServer:
 
         # You need to have at least acceleration, so you want 2 timesteps of prior data, e.g. [0, 1],
         # so that you can immediately start incremental inference from the 3rd timestep onwards.
-        init_timestep = 1
+        init_timestep = 2
 
         eval_scene = eval_env.scenes[scene_idx]
         online_env = create_online_env(eval_env, hyperparams, scene_idx, init_timestep)
@@ -311,24 +314,25 @@ class PredictionServer:
                 maps = get_maps_for_input(input_dict, eval_scene, hyperparams)
             # print(maps)
 
-            robot_present_and_future = None
-            if eval_scene.robot is not None and hyperparams['incl_robot_node']:
-                robot_present_and_future = eval_scene.robot.get(np.array([timestep,
-                                                                        timestep + hyperparams['prediction_horizon']]),
-                                                                hyperparams['state'][eval_scene.robot.type],
-                                                                padding=0.0)
-                robot_present_and_future = np.stack([robot_present_and_future, robot_present_and_future], axis=0)
-                # robot_present_and_future += adjustment
+            # robot_present_and_future = None
+            # if eval_scene.robot is not None and hyperparams['incl_robot_node']:
+            #     robot_present_and_future = eval_scene.robot.get(np.array([timestep,
+            #                                                             timestep + hyperparams['prediction_horizon']]),
+            #                                                     hyperparams['state'][eval_scene.robot.type],
+            #                                                     padding=0.0)
+            #     robot_present_and_future = np.stack([robot_present_and_future, robot_present_and_future], axis=0)
+            #     # robot_present_and_future += adjustment
 
             start = time.time()
             with warnings.catch_warnings():
                 warnings.simplefilter('ignore') # prevent deluge of UserWarning from torch's rrn.py
                 dists, preds = trajectron.incremental_forward(input_dict,
                                                             maps,
-                                                            prediction_horizon=125, # TODO: make variable
-                                                            num_samples=5, # TODO: make variable
-                                                            robot_present_and_future=robot_present_and_future,
-                                                            full_dist=True)
+                                                            prediction_horizon=self.config.prediction_horizon, # TODO: make variable
+                                                            num_samples=self.config.num_samples, # TODO: make variable
+                                                            full_dist=self.config.full_dist,
+                                                            gmm_mode=self.config.gmm_mode,
+                                                            z_mode=self.config.z_mode)
             end = time.time()
             logger.debug("took %.2f s (= %.2f Hz) w/ %d nodes and %d edges" % (end - start,
                                                                             1. / (end - start), len(trajectron.nodes),
@@ -387,6 +391,10 @@ class PredictionServer:
                 logger.info(f"Frame prediction: {len(trajectron.nodes)} nodes & {trajectron.scene_graph.get_num_edges()} edges. Trajectron: {end - start}s")
             else:
                 logger.info(f"Total frame delay = {time.time()-frame.time}s ({len(trajectron.nodes)} nodes & {trajectron.scene_graph.get_num_edges()} edges. Trajectron: {end - start}s)")
+            
+            if self.config.smooth_predictions:
+                frame = self.smoother.smooth_frame_predictions(frame)
+
             self.prediction_socket.send_pyobj(frame)
         logger.info('Stopping')
 

trap/renderer.py

@@ -99,7 +99,7 @@ class Renderer:
             if first_time is None:
                     first_time = frame.time
 
-            decorate_frame(frame, prediction_frame, first_time)
+            decorate_frame(frame, prediction_frame, first_time, self.config)
 
             img_path = (self.config.output_dir / f"{i:05d}.png").resolve()
 
@@ -107,9 +107,9 @@ class Renderer:
             
             logger.debug(f"write frame {frame.time - first_time:.3f}s")
             if self.out_writer:
-                self.out_writer.write(img)
+                self.out_writer.write(frame.img)
             if self.streaming_process:
-                self.streaming_process.stdin.write(img.tobytes())
+                self.streaming_process.stdin.write(frame.img.tobytes())
         logger.info('Stopping')
 
         if i>2:
@@ -121,7 +121,25 @@ class Renderer:
                 # oddly wrapped, because both close and release() take time.
                 self.streaming_process.wait()
 
-def decorate_frame(frame: Frame, prediction_frame: Frame, first_time) -> np.array:
+# colorset = itertools.product([0,255], repeat=3) # but remove white
+colorset = [(0, 0, 0),
+ (0, 0, 255),
+ (0, 255, 0),
+ (0, 255, 255),
+ (255, 0, 0),
+ (255, 0, 255),
+ (255, 255, 0)
+ ]
+
+
+def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, config: Namespace) -> np.array:
+    frame.img
+
+    overlay = np.zeros(frame.img.shape, np.uint8)
+    # Fill image with red color(set each pixel to red)
+    overlay[:] = (128, 0, 128)
+
+    frame.img = cv2.addWeighted(frame.img, .5, overlay, .5, 0)
     img = frame.img
 
     # all not working:
@@ -132,9 +150,10 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time) -> np.arra
     #     new_H = S * self.H * np.linalg.inv(S)
     #     warpedFrame = cv2.warpPerspective(img, new_H, (1000,1000))
     #     cv2.imwrite(str(self.config.output_dir / "orig.png"), warpedFrame)
+    cv2.rectangle(img, (0,0), (img.shape[1],25), (0,0,0), -1)
     
     if not prediction_frame:
-        cv2.putText(img, f"Waiting for prediction...", (20,50), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
+        cv2.putText(img, f"Waiting for prediction...", (20,20), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
         # continue
     else:
         inv_H = np.linalg.pinv(prediction_frame.H)
@@ -151,46 +170,69 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time) -> np.arra
             for ci in range(1, len(coords)):
                 start = [int(p) for p in coords[ci-1]]
                 end = [int(p) for p in coords[ci]]
-                color = (255,255,255) if confirmations[ci] else (100,100,100)
-                cv2.line(img, start, end, color, 2, lineType=cv2.LINE_AA)
+                # color = (255,255,255) if confirmations[ci] else (100,100,100)
+                color = [100+155*ci/len(coords)]*3
+                cv2.line(img, start, end, color, 1, lineType=cv2.LINE_AA)
+                cv2.circle(img, end, 2, color, lineType=cv2.LINE_AA)
 
             if not track.predictions or not len(track.predictions):
                 continue
 
+            color = colorset[track_id % len(colorset)]
+
             for pred_i, pred in enumerate(track.predictions):
-                pred_coords = cv2.perspectiveTransform(np.array([pred]), inv_H)[0]
-                color = (0,0,255) if pred_i else (100,100,100)
-                for ci in range(1, len(pred_coords)):
+                pred_coords = cv2.perspectiveTransform(np.array([pred]), inv_H)[0].tolist()
+                # color = (128,0,128) if pred_i else (128,128,0)
+                
+                for ci in range(0, len(pred_coords)):
+                    if ci == 0:
+                        start = [int(p) for p in coords[-1]]
+                        # start = [0,0]?
+                        # print(start)
+                    else:
                         start = [int(p) for p in pred_coords[ci-1]]
                     end = [int(p) for p in pred_coords[ci]]
-                    cv2.line(img, start, end, color, 1, lineType=cv2.LINE_AA)
+                    cv2.line(img, start, end, color, 2, lineType=cv2.LINE_AA)
+                    cv2.circle(img, end, 2, color, 1, lineType=cv2.LINE_AA)
             
         for track_id, track in prediction_frame.tracks.items():
             # draw tracker marker and track id last so it lies over the trajectories
             # this goes is a second loop so it overlays over _all_ trajectories
             # coords = cv2.perspectiveTransform(np.array([[track.history[-1].get_foot_coords()]]), self.inv_H)[0]
             coords = track.history[-1].get_foot_coords()
+            color = colorset[track_id % len(colorset)]
 
             center = [int(p) for p in coords]
-            cv2.circle(img, center, 5, (0,255,0))
+            cv2.circle(img, center, 6, (255,255,255), thickness=3)
             (l, t, r, b) = track.history[-1].to_ltrb()
             p1 = (l, t)
             p2 = (r, b)
-            cv2.rectangle(img, p1, p2, (255,0,0), 1)
-            cv2.putText(img, f"{track_id} ({(track.history[-1].conf or 0):.2f})", (center[0]+8, center[1]), cv2.FONT_HERSHEY_SIMPLEX, fontScale=.7, thickness=2, color=(0,255,0), lineType=cv2.LINE_AA)
+            # cv2.rectangle(img, p1, p2, (255,0,0), 1)
+            cv2.putText(img, f"{track_id} ({(track.history[-1].conf or 0):.2f})", (center[0]+8, center[1]), cv2.FONT_HERSHEY_SIMPLEX, fontScale=.7, thickness=1, color=color, lineType=cv2.LINE_AA)
 
 
-    cv2.putText(img, f"{frame.index:06d}", (20,50), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
-    cv2.putText(img, f"{frame.time - first_time:.3f}s", (120,50), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
+    base_color = (255,)*3
+    info_color = (255,255,0)
+       
+    cv2.putText(img, f"{frame.index:06d}", (20,17), cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
+    cv2.putText(img, f"{frame.time - first_time:.3f}s", (120,17), cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
 
     if prediction_frame:
         # render Δt and Δ frames
-        cv2.putText(img, f"{prediction_frame.index - frame.index}", (90,50), cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 1)
-        cv2.putText(img, f"{prediction_frame.time - time.time():.2f}s", (200,50), cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 1)
-        cv2.putText(img, f"{len(prediction_frame.tracks)} tracks", (500,50), cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 1)
-        cv2.putText(img, f"h: {np.average([len(t.history or []) for t in prediction_frame.tracks.values()])}", (580, 50),  cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 1)
-        cv2.putText(img, f"ph: {np.average([len(t.predictor_history or []) for t in prediction_frame.tracks.values()])}", (660, 50),  cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 1)
-        cv2.putText(img, f"p: {np.average([len(t.predictions or []) for t in prediction_frame.tracks.values()])}", (740, 50),  cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 1)
+        cv2.putText(img, f"{prediction_frame.index - frame.index}", (90,17), cv2.FONT_HERSHEY_PLAIN, 1, info_color, 1)
+        cv2.putText(img, f"{prediction_frame.time - time.time():.2f}s", (200,17), cv2.FONT_HERSHEY_PLAIN, 1, info_color, 1)
+        cv2.putText(img, f"{len(prediction_frame.tracks)} tracks", (500,17), cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
+        cv2.putText(img, f"h: {np.average([len(t.history or []) for t in prediction_frame.tracks.values()])}", (580,17),  cv2.FONT_HERSHEY_PLAIN, 1, info_color, 1)
+        cv2.putText(img, f"ph: {np.average([len(t.predictor_history or []) for t in prediction_frame.tracks.values()])}", (660,17),  cv2.FONT_HERSHEY_PLAIN, 1, info_color, 1)
+        cv2.putText(img, f"p: {np.average([len(t.predictions or []) for t in prediction_frame.tracks.values()])}", (740,17),  cv2.FONT_HERSHEY_PLAIN, 1, info_color, 1)
+    
+    options = []
+    for option in ['prediction_horizon','num_samples','full_dist','gmm_mode','z_mode', 'model_dir']:
+        options.append(f"{option}: {config.__dict__[option]}")
+    
+
+    cv2.putText(img, options.pop(-1), (20,img.shape[0]-30),  cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
+    cv2.putText(img, " | ".join(options), (20,img.shape[0]-10),  cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
 
     return img
 

trap/tracker.py

@@ -24,6 +24,10 @@ from ultralytics.engine.results import Results as YOLOResult
 
 from trap.frame_emitter import DetectionState, Frame, Detection, Track
 
+
+from tsmoothie.smoother import KalmanSmoother, ConvolutionSmoother
+import tsmoothie.smoother
+
 # Detection = [int, int, int, int, float, int]
 # Detections = [Detection]
 
@@ -103,6 +107,13 @@ class Tracker:
 
         self.H = np.loadtxt(self.config.homography, delimiter=',')
 
+        if self.config.smooth_tracks:
+            logger.info("Smoother enabled")
+            self.smoother = Smoother()
+        else:
+            logger.info("Smoother Disabled (enable with --smooth-tracks)")
+
+
         logger.debug("Set up tracker")
         
 
@@ -160,7 +171,7 @@ class Tracker:
 
 
             if self.config.detector == DETECTOR_YOLOv8:
-                detections: [Detection]  = self._yolov8_track(frame.img)
+                detections: [Detection]  = self._yolov8_track(frame)
             else :
                 detections: [Detection] = self._resnet_track(frame.img, scale = 1)
             
@@ -199,6 +210,9 @@ class Tracker:
             #     self.trajectory_socket.send_string(json.dumps(trajectories))
             # else:
             #     self.trajectory_socket.send(pickle.dumps(frame))
+            if self.config.smooth_tracks:
+                frame = self.smoother.smooth_frame_tracks(frame)
+
             self.trajectory_socket.send_pyobj(frame)
                 
             current_time = time.time()
@@ -249,12 +263,12 @@ class Tracker:
                     
         logger.info('Stopping')
 
-    def _yolov8_track(self, img) -> [Detection]:
-        results: [YOLOResult] = self.model.track(img, persist=True)
+    def _yolov8_track(self, frame: Frame,) -> [Detection]:
+        results: [YOLOResult] = self.model.track(frame.img, persist=True, tracker="bytetrack.yaml", verbose=False)
         if results[0].boxes is None or results[0].boxes.id is None:
             # work around https://github.com/ultralytics/ultralytics/issues/5968
             return []
-        return [Detection(track_id, *bbox) for bbox, track_id in zip(results[0].boxes.xywh.cpu(), results[0].boxes.id.int().cpu().tolist())]
+        return [Detection(track_id, bbox[0]-.5*bbox[2], bbox[1]-.5*bbox[3], bbox[2], bbox[3], 1, DetectionState.Confirmed, frame.index) for bbox, track_id in zip(results[0].boxes.xywh.cpu(), results[0].boxes.id.int().cpu().tolist())]
         
     def _resnet_track(self, img, scale: float = 1) -> [Detection]:
         if scale != 1:
@@ -304,3 +318,55 @@ class Tracker:
 def run_tracker(config: Namespace, is_running: Event):
     router = Tracker(config, is_running)
     router.track()
+
+
+
+class Smoother:
+
+    def __init__(self, window_len=2):
+        self.smoother = ConvolutionSmoother(window_len=window_len, window_type='ones', copy=None)
+        
+    
+    def smooth_frame_tracks(self, frame: Frame) -> Frame:
+        new_tracks = []
+        for track in frame.tracks.values():
+            ls = [d.l for d in track.history]
+            ts = [d.t for d in track.history]
+            ws = [d.w for d in track.history]
+            hs = [d.h for d in track.history]
+            self.smoother.smooth(ls)
+            ls = self.smoother.smooth_data[0]
+            self.smoother.smooth(ts)
+            ts = self.smoother.smooth_data[0]
+            self.smoother.smooth(ws)
+            ws = self.smoother.smooth_data[0]
+            self.smoother.smooth(hs)
+            hs = self.smoother.smooth_data[0]
+            new_history = [Detection(d.track_id, l, t, w, h, d.conf, d.state, d.frame_nr) for l, t, w, h, d in zip(ls,ts,ws,hs, track.history)]
+            new_track = Track(track.track_id, new_history, track.predictor_history, track.predictions)
+            new_tracks.append(new_track)
+        frame.tracks = {t.track_id: t for t in new_tracks}
+        return frame
+    
+    def smooth_frame_predictions(self, frame) -> Frame:
+        
+        for track in frame.tracks.values():
+            new_predictions = []
+            if not track.predictions:
+                continue
+            
+            for prediction in track.predictions:
+                xs = [d[0] for d in prediction]
+                ys = [d[1] for d in prediction]
+            
+                self.smoother.smooth(xs)
+                xs = self.smoother.smooth_data[0]
+                self.smoother.smooth(ys)
+                ys = self.smoother.smooth_data[0]
+                
+                smooth_prediction = [[x,y] for x, y in zip(xs, ys)]
+                
+                new_predictions.append(smooth_prediction)
+            track.predictions = new_predictions
+        
+        return frame