trap/config.py

@@ -152,31 +152,6 @@ 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.
 
@@ -217,10 +192,6 @@ 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",
@@ -243,9 +214,6 @@ 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,7 +42,6 @@ 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]
@@ -150,12 +149,6 @@ 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, Smoother
+from trap.tracker import Track
 
 logger = logging.getLogger("trap.prediction")
 
@@ -121,9 +121,6 @@ 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'')
@@ -191,7 +188,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 = 2
+        init_timestep = 1
 
         eval_scene = eval_env.scenes[scene_idx]
         online_env = create_online_env(eval_env, hyperparams, scene_idx, init_timestep)
@@ -314,25 +311,24 @@ class PredictionServer:
                 maps = get_maps_for_input(input_dict, eval_scene, hyperparams)
             # print(maps)
 
-            # robot_present_and_future = None
+            robot_present_and_future = None
-            # if eval_scene.robot is not None and hyperparams['incl_robot_node']:
+            if eval_scene.robot is not None and hyperparams['incl_robot_node']:
-            #     robot_present_and_future = eval_scene.robot.get(np.array([timestep,
+                robot_present_and_future = eval_scene.robot.get(np.array([timestep,
-            #                                                             timestep + hyperparams['prediction_horizon']]),
+                                                                        timestep + hyperparams['prediction_horizon']]),
-            #                                                     hyperparams['state'][eval_scene.robot.type],
+                                                                hyperparams['state'][eval_scene.robot.type],
-            #                                                     padding=0.0)
+                                                                padding=0.0)
-            #     robot_present_and_future = np.stack([robot_present_and_future, robot_present_and_future], axis=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 += 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=self.config.prediction_horizon, # TODO: make variable
+                                                            prediction_horizon=125, # TODO: make variable
-                                                            num_samples=self.config.num_samples, # TODO: make variable
+                                                            num_samples=5, # TODO: make variable
-                                                            full_dist=self.config.full_dist,
+                                                            robot_present_and_future=robot_present_and_future,
-                                                            gmm_mode=self.config.gmm_mode,
+                                                            full_dist=True)
-                                                            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),
@@ -391,10 +387,6 @@ 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, self.config)
+            decorate_frame(frame, prediction_frame, first_time)
 
             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(frame.img)
+                self.out_writer.write(img)
             if self.streaming_process:
-                self.streaming_process.stdin.write(frame.img.tobytes())
+                self.streaming_process.stdin.write(img.tobytes())
         logger.info('Stopping')
 
         if i>2:
@@ -121,25 +121,7 @@ class Renderer:
                 # oddly wrapped, because both close and release() take time.
                 self.streaming_process.wait()
 
-# colorset = itertools.product([0,255], repeat=3) # but remove white
+def decorate_frame(frame: Frame, prediction_frame: Frame, first_time) -> np.array:
-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:
@@ -150,10 +132,9 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, con
     #     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,20), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
+        cv2.putText(img, f"Waiting for prediction...", (20,50), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
         # continue
     else:
         inv_H = np.linalg.pinv(prediction_frame.H)
@@ -170,69 +151,46 @@ def decorate_frame(frame: Frame, prediction_frame: Frame, first_time: float, con
             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)
+                color = (255,255,255) if confirmations[ci] else (100,100,100)
-                color = [100+155*ci/len(coords)]*3
+                cv2.line(img, start, end, color, 2, lineType=cv2.LINE_AA)
-                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].tolist()
+                pred_coords = cv2.perspectiveTransform(np.array([pred]), inv_H)[0]
-                # color = (128,0,128) if pred_i else (128,128,0)
+                color = (0,0,255) if pred_i else (100,100,100)
-                
+                for ci in range(1, len(pred_coords)):
-                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, 2, lineType=cv2.LINE_AA)
+                    cv2.line(img, start, end, color, 1, 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, 6, (255,255,255), thickness=3)
+            cv2.circle(img, center, 5, (0,255,0))
             (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.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"{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)
 
        
-    base_color = (255,)*3
+    cv2.putText(img, f"{frame.index:06d}", (20,50), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
-    info_color = (255,255,0)
+    cv2.putText(img, f"{frame.time - first_time:.3f}s", (120,50), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,0), 1)
-       
-    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,17), cv2.FONT_HERSHEY_PLAIN, 1, info_color, 1)
+        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,17), cv2.FONT_HERSHEY_PLAIN, 1, info_color, 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,17), cv2.FONT_HERSHEY_PLAIN, 1, base_color, 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,17),  cv2.FONT_HERSHEY_PLAIN, 1, info_color, 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,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, 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,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, 50),  cv2.FONT_HERSHEY_PLAIN, 1, (0,0,255), 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,10 +24,6 @@ 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]
 
@@ -107,13 +103,6 @@ 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")
         
 
@@ -171,7 +160,7 @@ class Tracker:
 
 
             if self.config.detector == DETECTOR_YOLOv8:
-                detections: [Detection]  = self._yolov8_track(frame)
+                detections: [Detection]  = self._yolov8_track(frame.img)
             else :
                 detections: [Detection] = self._resnet_track(frame.img, scale = 1)
             
@@ -210,9 +199,6 @@ 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()
@@ -263,12 +249,12 @@ class Tracker:
                     
         logger.info('Stopping')
 
-    def _yolov8_track(self, frame: Frame,) -> [Detection]:
+    def _yolov8_track(self, img) -> [Detection]:
-        results: [YOLOResult] = self.model.track(frame.img, persist=True, tracker="bytetrack.yaml", verbose=False)
+        results: [YOLOResult] = self.model.track(img, persist=True)
         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[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())]
+        return [Detection(track_id, *bbox) 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:
@@ -318,55 +304,3 @@ 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