diff --git a/trap/cv_renderer.py b/trap/cv_renderer.py index 2d705e8..add1051 100644 --- a/trap/cv_renderer.py +++ b/trap/cv_renderer.py @@ -503,6 +503,21 @@ def decorate_frame(frame: Frame, tracker_frame: Frame, prediction_frame: Frame, anim_position = get_animation_position(track, frame) draw_track_predictions(img, track, int(track.track_id)+1, config.camera, convert_world_points_to_img_points, anim_position=anim_position) cv2.putText(img, f"{len(track.predictor_history) if track.predictor_history else 'none'}", to_point(track.history[0].get_foot_coords()), cv2.FONT_HERSHEY_COMPLEX, 1, (255,255,255), 1) + if prediction_frame.maps: + for i, m in enumerate(prediction_frame.maps): + map_img = np.ascontiguousarray(np.flipud(np.transpose(m[0], (2, 1, 0))*255), np.uint8) + cv2.circle(map_img, (10,50), 5, (255,0,0), 2) + + height, width, _ = map_img.shape + padding= 50 + y = img.shape[0] - padding - height + x = width*i + + if x+width > img.shape[1]: + break # stop drawing maps when there's a lot of them + + img[y:y+height,x:x+width] = map_img + base_color = (255,)*3 diff --git a/trap/frame_emitter.py b/trap/frame_emitter.py index 5f39fdc..4c722d4 100644 --- a/trap/frame_emitter.py +++ b/trap/frame_emitter.py @@ -406,6 +406,7 @@ class Frame: tracks: Optional[dict[str, Track]] = None H: Optional[np.array] = None camera: Optional[Camera] = None + maps: Optional[List[cv2.Mat]] = None def aslist(self) -> [dict]: return { t.track_id: @@ -421,7 +422,7 @@ class Frame: } def without_img(self): - return Frame(self.index, None, self.time, self.tracks, self.H, self.camera) + return Frame(self.index, None, self.time, self.tracks, self.H, self.camera, self.maps) def video_src_from_config(config) -> UrlOrPath: if config.video_loop: diff --git a/trap/prediction_server.py b/trap/prediction_server.py index 7f7e64c..3a45870 100644 --- a/trap/prediction_server.py +++ b/trap/prediction_server.py @@ -378,7 +378,7 @@ class PredictionServer: if hyperparams['use_map_encoding']: maps = get_maps_for_input(input_dict, eval_scene, hyperparams, device=self.config.eval_device) - print(maps) + # print(maps) # robot_present_and_future = None # if eval_scene.robot is not None and hyperparams['incl_robot_node']: @@ -479,6 +479,8 @@ class PredictionServer: if self.config.smooth_predictions: frame = self.smoother.smooth_frame_predictions(frame) + + frame.maps = list([m.cpu().numpy() for m in maps.values()]) if maps else None self.send_frame(frame)