render trajectron's cropped maps in cv renderer

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

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:

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)