minor fixes

ToyotaResearchInstitute · Feb 2, 2022 · ab0e80f · ab0e80f
1 parent 61d174d
commit ab0e80f
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/src/maad/dataset/maad_base_dataset.py b/src/maad/dataset/maad_base_dataset.py
@@ -488,6 +488,8 @@ def _get_single_item(self, metadata):
                 video_id of the video corresponding to the data item
             AUXILIARY_INFO_SUBJECT_ID: int
                 subject id of the subject corresponding to the data_item
+            AUXILIARY_INFO_FRAME_IDX: int
+                frame index of the corresponding data item
             AUXILIARY_INFO_FULL_SIZE_GAZE_0: numpy.array (L, 2)
                 Gaze points in full resolution
 
@@ -536,6 +538,7 @@ def _get_single_item(self, metadata):
         else:
             auxiliary_info = collections.OrderedDict()
             auxiliary_info[AUXILIARY_INFO_VIDEO_ID] = video_id
+            auxiliary_info[AUXILIARY_INFO_FRAME_IDX] = frame_idx
             auxiliary_info[AUXILIARY_INFO_SUBJECT_ID] = subject
             auxiliary_info[AUXILIARY_INFO_FULL_SIZE_GAZE_0] = full_size_gaze_points_array
 

diff --git a/src/maad/utils/maad_consts.py b/src/maad/utils/maad_consts.py
@@ -40,6 +40,7 @@
 # daze Dataset specific field names.
 AUXILIARY_INFO_VIDEO_ID = "video_id"
 AUXILIARY_INFO_SUBJECT_ID = "subject_id"
+AUXILIARY_INFO_FRAME_IDX = "frame_id"
 AUXILIARY_INFO_FULL_SIZE_GAZE_0 = AUXILIARY_INFO_FULL_SIZE_GAZE_TEMPLATE.format(0)
 
 # enum for determining whether side channel gaze be used during inference.

diff --git a/src/scripts/generate_saliency_maps_using_model.py b/src/scripts/generate_saliency_maps_using_model.py
@@ -10,6 +10,7 @@
 
 from maad.model.gaze_transform import GazeTransform
 from maad.model.gaze_corruption import GazeCorruption
+from maad.utils.maad_consts import *
 
 # (NOTE) When running this script please make sure that you are using --batch_size = 1
 
@@ -208,10 +209,9 @@ def convert_heatmap_to_image(
                 for t in t_list:
                     # Weird ordering due to the "list nature of aux_info_list"
                     print("Time step ", t)
-                    frame_idx = aux_info_list[t]["road_id"][1][b].numpy().item()
-                    # for example, 'DREYEVE_DATA/06/video_garmin.avi'
-                    video_id_path = aux_info_list[t]["road_id"][0][b]
-                    video_id_num = os.path.split(os.path.split(video_id_path)[0])[1]  # for example, '06'
+                    frame_idx = aux_info_list[t][AUXILIARY_INFO_FRAME_IDX][b].numpy().item()
+                    # for example, '06'
+                    video_id_num = "{0:02d}".format(aux_info_list[t][AUXILIARY_INFO_VIDEO_ID][b])
                     gaze_hmap_with_gaze = (
                         predicted_gaze_with_gaze["gaze_density_map"][b, t, 0, :, :].cpu().detach().numpy()
                     )