supports random generation for multichoice dataset

Author: Haozhe Qi

llava/action/dataset.py

@@ -6,7 +6,7 @@
 import torch
 import decord
 from pathlib import Path
-from llava.action.utils import  AvionMultiChoiceGenerator,  avion_video_loader, EK100_frame_loader
+from llava.action.utils import  AvionMultiChoiceGenerator,  RandomMultiChoiceGenerator, avion_video_loader, EK100_frame_loader
 from llava.action.prediction_analysis import PredictionAnalysis
 import torch.distributed as dist
 
@@ -147,6 +147,7 @@ def __init__(
         verb_maps = None,
         noun_maps = None,
         eval_result_folder = None,
+        gen_type = 'action_model',
         action_representation = 'GT_random_narration_cut',
         mapping_vn2narration = None,
         avion_predictions = None,
@@ -175,14 +176,19 @@ def __init__(
         self.labels = labels
         self.topk_predictions = topk_predictions
         self.ann_root = Path(metadata).parent
-        self.mc_generator = AvionMultiChoiceGenerator(self.ann_root)
+        self.gen_type = gen_type
+        if gen_type == 'action_model':
+            self.mc_generator = AvionMultiChoiceGenerator(self.ann_root)
+        elif gen_type == 'random':
+            self.mc_generator = RandomMultiChoiceGenerator(self.ann_root)
         self.rank = dist.get_rank()
         self.prediction_analysis = PredictionAnalysis(rank = self.rank, save_folder = eval_result_folder)
         self.action_representation = action_representation
         self.n_narrations = n_narrations
         self.mapping_vn2narration = mapping_vn2narration
         self.avion_predictions = avion_predictions
 
+        
     def __getitem__(self, i):
         frames, label, time_meta = self.get_raw_item(
             i, is_training=self.is_training,
@@ -205,19 +211,31 @@ def __getitem__(self, i):
             frames = self.transform(frames)
         narration = self.samples[i][4]
         avion_preds = self.avion_predictions[str(i)]['predictions']
-
-        data = self.mc_generator.generate_multi_choice(label, 
-                                                        avion_preds,                                                       
-                                                        narration,
-                                                        self.topk_predictions, 
-                                                        self.action_representation,
-                                                        self.n_narrations,
-                                                        self.labels,
-                                                        self.mapping_vn2narration,                                                        
-                                                        self.verb_maps, 
-                                                        self.noun_maps,
-                                                        benchmark_testing = self.eval_args.benchmark_testing,
-                                                        is_train = False) # note we only use this dataset for evaluation for now.
-
+        if self.gen_type =='action_model':
+            data = self.mc_generator.generate_multi_choice(label, 
+                                                            avion_preds,                                                       
+                                                            narration,
+                                                            self.topk_predictions, 
+                                                            self.action_representation,
+                                                            self.n_narrations,
+                                                            self.labels,
+                                                            self.mapping_vn2narration,                                                        
+                                                            self.verb_maps, 
+                                                            self.noun_maps,
+                                                            benchmark_testing = self.eval_args.benchmark_testing,
+                                                            is_train = False) # note we only use this dataset for evaluation for now.
+        else:
+            data = self.mc_generator.generate_multi_choice(label, 
+                                                            narration,
+                                                            self.topk_predictions, 
+                                                            self.action_representation,
+                                                            self.n_narrations,
+                                                            self.labels,
+                                                            self.mapping_vn2narration,                                                        
+                                                            self.verb_maps, 
+                                                            self.noun_maps,
+                                                            benchmark_testing = self.eval_args.benchmark_testing,
+                                                            is_train = False) # no
+            
 
         return frames, data, time_meta, i

llava/action/ek_eval.py

@@ -131,6 +131,7 @@ def get_args_parser():
     parser.add_argument("--perspective", default = "first_person", type = str)
     parser.add_argument('--benchmark_testing', action='store_true', default = False)
     parser.add_argument('--include_time_instruction', action='store_true', default = False)
+    parser.add_argument('--gen_type', type = str, default = 'action_model') # action_model, random
     return parser
 
 def prepare_llava(pretrained):
@@ -191,7 +192,7 @@ def ensemble_llava_evaluation(
     # shuffle the options
     options = mc_data['options'][0]
     letters = mc_data['valid_letters']
-    avion_pred = mc_data['avion_pred']
+    avion_pred = mc_data.get('avion_pred', None)
     # each option was in the format of {letter}. {answer}
     preds = []
     for _ in range(ensemble_k):
@@ -283,6 +284,7 @@ def evaluate_on_EK100(eval_args,
                 mapping_vn2narration = mapping_vn2narration,
                 avion_predictions = predictions if eval_args.action_predictions else None,
                 n_narrations = eval_args.n_narrations,
+                gen_type = eval_args.gen_type
             )
 
     def collate_fn(batch):
@@ -371,7 +373,7 @@ def collate_fn(batch):
             local_running_corrects = torch.tensor(0.0, device=device)
             local_total_samples = torch.tensor(0.0, device=device)            
 
-            if eval_args.action_predictions:
+            if eval_args.action_predictions and eval_args.gen_type == 'action_model':
                 avion_pred = mc_data['avion_pred']
                 if gt_name == avion_pred:               
                     local_avion_correct.add_(1)
@@ -420,7 +422,7 @@ def collate_fn(batch):
             val_dataset.prediction_analysis.log(global_index,
                                                 llava_pred,
                                                 gt_name,
-                                                mc_data['all_avion_preds'],
+                                                mc_data.get('all_avion_preds', None),
                                                 time_meta['start_second'],
                                                 time_meta['end_second'],
                                                 time_meta['vid_path'],

llava/action/utils.py

@@ -425,7 +425,6 @@ def generate_multi_choice(self,
         randomly pick k-1 letters from vn_list
 
         """        
-        
         if is_train:
             return self.train_generate(gt_vn, narration, k, action_representation, n_narrations, labels, mapping_vn2narration, verb_maps, noun_maps)
         else:
@@ -547,14 +546,13 @@ def test_generate(self,
         answer_ids = action_model_predictions[:k]
 
         if benchmark_testing:
-            print ("am i here")
             # if we are testing on benchmark, we need to ensure that the gt_vn is in the top k predictions
             # if not, we remove the last prediction and add the gt_vn
             if gt_vn not in answer_ids:
                 answer_ids.pop()
                 answer_ids.append(gt_vn)
         else:
-           print ("am i not here")
+           pass
 
         answers = []
         for answer_id in answer_ids:

llava/train/train.py

@@ -206,6 +206,7 @@ class EK100EvalArguments:
     pseudo_folder: str = ""
     benchmark_testing: bool = False
     include_time_instruction: bool = False
+    gen_type: str = 'action_model'
 
 def maybe_zero_3(param, ignore_status=False, name=None):
     from deepspeed import zero