Update visualise.py

Author: HRHLALALA

code/visualise.py

@@ -1,3 +1,4 @@
+import math
 import pickle
 import re
 from collections import defaultdict
@@ -8,6 +9,39 @@
 import numpy as np
 import cv2
 
+activity2id = {
+    "BG": 0,  # background
+    "activity_walking": 1,
+    "activity_standing": 2,
+    "activity_carrying": 3,
+    "activity_gesturing": 4,
+    "Closing": 5,
+    "Opening": 6,
+    "Interacts": 7,
+    "Exiting": 8,
+    "Entering": 9,
+    "Talking": 10,
+    "Transport_HeavyCarry": 11,
+    "Unloading": 12,
+    "Pull": 13,
+    "Loading": 14,
+    "Open_Trunk": 15,
+    "Closing_Trunk": 16,
+    "Riding": 17,
+    "specialized_texting_phone": 18,
+    "Person_Person_Interaction": 19,
+    "specialized_talking_phone": 20,
+    "activity_running": 21,
+    "PickUp": 22,
+    "specialized_using_tool": 23,
+    "SetDown": 24,
+    "activity_crouching": 25,
+    "activity_sitting": 26,
+    "Object_Transfer": 27,
+    "Push": 28,
+    "PickUp_Person_Vehicle": 29,
+}
+
 
 def get_person_box_at_frame(tracking_results, target_frame_idx):
     """Get all the box from the tracking result at frameIdx."""
@@ -86,12 +120,11 @@ def visualise(args):
     with open(args.pred_path, "rb") as f:
         pred_data = pickle.load(f)
         pred_data = {
-            k: v for k, v in pred_data.items() if k in ['obs_list', 'pred_gt_list', 'pred_list']
+            k: v for k, v in pred_data.items() if k in ['seq_ids', 'obs_list', 'pred_gt_list', 'pred_list', 'pred_act']
         }
 
     video_dict = get_video_dict(pred_data)
-
-    del pred_data['act_pred_list']
+    id2activity = {v:k for k,v in activity2id.items()}
     for video_name in video_dict:
         video_file = os.path.join(args.video_base, video_name + '.mp4')
 
@@ -117,7 +150,7 @@ def visualise(args):
             #
             #
             #     print(seq_id, video_meta)
-        cur_frame = -1
+        cur_frame = - 8 * args.drop_frame
         next_signal = False
         while cur_frame < video_meta["frame_count"]:
             cur_frame += 1
@@ -131,15 +164,21 @@ def visualise(args):
 
             for idx in filter_pred(video_pred_data, video_name + f'_{cur_frame}_+'):
                 seq_id = video_pred_data['seq_ids'][idx]
-                obs_traj = video_pred_data['obs_traj_list'][idx]
-                pred_traj = video_pred_data['pred_traj_list'][idx]
-                gt_traj = video_pred_data['gt_traj_list'][idx]
-
+                obs_traj = video_pred_data['obs_list'][idx]
+                pred_traj = video_pred_data['pred_list'][idx]
+                gt_traj = video_pred_data['pred_gt_list'][idx]
+                # pred_act =sigmoid(np.array( video_pred_data['pred_act'][idx] * 20 ))
+                pred_act = video_pred_data['pred_act'][idx]
                 frame = cv2.resize(frame, (1920, 1080))
                 frame = plot_traj(frame, obs_traj, (255, 0, 0))
                 frame = plot_traj(frame, pred_traj, (0, 255, 0))
                 frame = plot_traj(frame, gt_traj, (0, 0, 255))
+                act_ids = np.nonzero( pred_act > 0.7)[0]
+                print(act_ids)
+                msg = ','.join([id2activity[id] for id in act_ids])
 
+                frame = cv2.putText(frame,msg, (int(obs_traj[-1,0]), int(obs_traj[-1,1])),cv2.FONT_HERSHEY_SIMPLEX,
+                            1, (255, 0, 0), 1, cv2.LINE_AA )
             # center = obs_traj[-1].astype("int")
             # frame = frame[center[1] - 200: center[1] + 200, center[0] - 200: center[0] + 200]
             frame = cv2.putText(frame, video_name + ' ' + 'frame ' + str(cur_frame), (30,30), cv2.FONT_HERSHEY_SIMPLEX,