several fixes to temporal dpo

Author: Ye Shaokai

.vscode/launch.json

@@ -250,7 +250,7 @@
 //                 "--model_name_or_path", "lmms-lab/llava-onevision-qwen2-0.5b-ov",
 //                 "--version", "qwen_1_5",
 //                 "--data_path", "scripts/train/simple_avion_top5_gt_and_direct.yaml",
-//                 "--video_folder", "/data/shaokai/",
+//                 "--video_folder", "/data/shaokai/EK100_512/",
 //                 "--mm_tunable_parts", "mm_vision_tower,mm_mlp_adapter,mm_language_model",
 //                 "--mm_vision_tower_lr", "2e-6",
 //                 "--vision_tower", "google/siglip-so400m-patch14-384",
@@ -288,7 +288,7 @@
 //                 "--torch_compile_backend", "inductor",
 //                 "--dataloader_drop_last", "True",
 //                 "--frames_upbound", "4",
-//                 "--root", "/data/shaokai/EK100",
+//                 "--root", "/data/shaokai/EK100_512/EK100",
 //                 "--action_predictions", "/data/shaokai/AVION_PREDS/avion_pred_ids_val.json",
 //                 "--val_metadata", "/data/shaokai/epic-kitchens-100-annotations/EPIC_100_validation.csv",
 //                 "--llava_num_frames", "4",
@@ -345,8 +345,9 @@
                     "--action_predictions","/data/shaokai/TIM_PREDS/tim_pred_ids_val.json",
                     "--action_representation", "official_key",
                     "--topk_predictions", "5",
-                    "--test_type", "base",
-                    "--output_dir", "test_0.5b_direct",                    
+                    "--test_type", "temporal_cot",
+                    "--output_dir", "test_0.5b_direct",  
+                    "--learn_neighbor_actions"                  
                 ],
                 "console": "integratedTerminal",
                 "justMyCode": false,

llava/action/dataset.py

@@ -221,3 +221,99 @@ def __getitem__(self, i):
 
         return frames, data, time_meta, i
 
+
+
+
+class VideoTemporalMultiChoiceDataset(VideoCaptionDatasetBase):
+    def __init__(
+        self, dataset, root, metadata, transform=None,
+        is_training=True, label_mapping=None,
+        num_clips=1,
+        chunk_len=300,
+        clip_length=32, clip_stride=2,
+        threads=1,
+        fast_rrc=False,
+        rrc_params=(224, (0.5, 1.0)),
+        fast_rcc=False,
+        rcc_params=(224,),
+        sparse_sample=False,
+        labels = None,
+        is_trimmed=True,
+        eval_args = None,
+        topk_predictions = 5,
+        verb_maps = None,
+        noun_maps = None,
+        eval_result_folder = None,
+        action_representation = 'GT_random_narration',
+        mapping_vn2narration = None,
+        avion_predictions = None,
+        n_narrations = -1,
+    ):
+        super().__init__(dataset, root, metadata, is_trimmed=is_trimmed)
+
+        self.transform = transform
+        self.is_training = is_training
+        self.label_mapping = label_mapping
+        self.num_clips = num_clips
+        self.chunk_len = chunk_len
+        self.clip_length = clip_length
+        self.clip_stride = clip_stride
+        self.threads = threads
+        self.fast_rrc = fast_rrc
+        self.rrc_params = rrc_params
+        self.fast_rcc = fast_rcc
+        self.rcc_params = rcc_params
+        self.sparse_sample = sparse_sample
+        self.eval_args = eval_args
+        self.verb_maps = verb_maps
+        self.noun_maps = noun_maps
+        self.vn_list = list(self.label_mapping.keys())        
+
+        self.labels = labels
+        self.topk_predictions = topk_predictions
+        self.ann_root = Path(metadata).parent
+        self.mc_generator = AvionMultiChoiceGenerator(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,
+            chunk_len=self.chunk_len,
+            num_clips=self.num_clips,
+            clip_length=self.clip_length,
+            clip_stride=self.clip_stride,
+            threads=self.threads,
+            fast_rrc=self.fast_rrc,
+            rrc_params=self.rrc_params,
+            fast_rcc=self.fast_rcc,
+            rcc_params=self.rcc_params,
+            sparse_sample=self.sparse_sample,
+        )
+
+        # for llava-video to work, we also need time meta data.
+
+        # apply transformation
+        if self.transform is not None:
+            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,
+                                                        is_train = False) # note we only use this dataset for evaluation for now.
+
+       
+        return frames, data, time_meta, i

llava/action/ek_eval.py

@@ -14,7 +14,7 @@
 from llava.action.utils import generate_label_map,  match_answer
 from collections import Counter 
 import torch.distributed as dist
-from llava.action.dataset import VideoMultiChoiceDataset
+from llava.action.dataset import VideoMultiChoiceDataset, VideoTemporalMultiChoiceDataset
 import torchvision.io as io
 import re
 
@@ -124,8 +124,9 @@ def get_args_parser():
                                    'random_narration_cut', 'top1_narration_cut', 'topk_narration_cut_key',
                                    'GT_key', 'GT_random_narration', 'GT_random_narration_cut', 'gpt_narration'])
     parser.add_argument('--n_narrations', default = -1, type = int)
-    parser.add_argument('--test_type', default = 'base', type = str, choices = ['caption', 'base', 'caption_then_answer', 'direct_narration'])
+    parser.add_argument('--test_type', default = 'base', type = str, choices = ['caption', 'base', 'temporal_cot', 'caption_then_answer', 'direct_narration'])
     parser.add_argument('--learn_neighbor_actions', action='store_true', default = False)
+    parser.add_argument('--pseudo_folder', default = None, type = str)
     parser.add_argument('--output_dir', default = None, type = str)
     return parser
 
@@ -253,7 +254,7 @@ def evaluate_on_EK100(eval_args,
     if eval_args.action_predictions:
         with open(eval_args.action_predictions, 'r') as f:
             predictions = json.load(f) 
-
+    
     val_dataset = VideoMultiChoiceDataset(
                 eval_args.dataset, eval_args.root, eval_args.val_metadata, val_transform_gpu,
                 is_training=False, label_mapping=mapping_vn2act,
@@ -332,12 +333,11 @@ def collate_fn(batch):
         os.makedirs('debug_and_vis', exist_ok = True)
 
 
-    uid_pad_dict = None
     lookup_table = None
     meta_data = None
     if eval_args.learn_neighbor_actions:
         from llava.action.generate_interval_pred import  get_lookup_dict
-        lookup_table = get_lookup_dict(eval_args.val_metadata)
+        lookup_table = get_lookup_dict(eval_args.val_metadata, test_type = eval_args.test_type, pseudo_folder = eval_args.pseudo_folder)
 
 
     for idx, (frames, mc_data, time_meta, global_index) in tqdm(enumerate(val_dataloader)):                  

llava/action/generate_interval_pred.py

@@ -103,16 +103,38 @@ def build_uid_pad_dict(ann_file,
     return uid_to_neighbors
 
 
+def get_pseudo_dict(pseudo_folder,  delta = 3):
+    import glob
+    
+
+    files = glob.glob(os.path.join(pseudo_folder, 'prediction*.json'))
+    
+    pseudo_data = {}
+    ret = {}
+    for file in files:
+        with open(file, 'r') as f:
+            pseudo_data.update(json.load(f))
+    for k,v in pseudo_data.items():
+        start_timestamp = round(float(v['start_second']),2)
+        end_timestamp = round(float(v['end_second']), 2)
+        vid = v['vid_path'].replace('/', '-')
+        uid = f"{vid}_{start_timestamp}_{end_timestamp}"
+        ret[uid] = v['llava_pred']
+            
+    assert len(ret) == len(pseudo_data)
+    return ret
 
-def get_lookup_dict(ann_file, delta = 3):
+def get_lookup_dict(ann_file, test_type = 'base', delta = 3, pseudo_folder = None):
 
     vid_to_intervals, vid_to_gt_narration, _ = get_annotated_intervals(ann_file)
     table = {}
 
+    pseudo_dict = None
+    if test_type == 'temporal_cot':
+        pseudo_dict = get_pseudo_dict(pseudo_folder)
+    
     for vid, intervals in vid_to_intervals.items():
-        
-        #sorted_intervals = sorted(intervals, key=lambda x: x[1])
-        
+                
         sorted_indices = sorted(range(len(intervals)), key=lambda i: intervals[i][1])
 
         sorted_intervals = [intervals[i] for i in sorted_indices]
@@ -136,10 +158,14 @@ def get_lookup_dict(ann_file, delta = 3):
                 uid2 = f"{id}_{round(start_times[i+1],2)}_{round(end_times[i+1],2)}"
                 uid3 = f"{id}_{round(start_times[i+2],2)}_{round(end_times[i+2],2)}"
 
-                
-                narration1 = sorted_narrations[i]
-                narration2 = sorted_narrations[i+1]
-                narration3 = sorted_narrations[i+2]
+                if test_type == 'base':
+                    narration1 = sorted_narrations[i]
+                    narration2 = sorted_narrations[i+1]
+                    narration3 = sorted_narrations[i+2]
+                elif test_type == 'temporal_cot':
+                    narration1 = pseudo_dict[uid1]
+                    narration2 = pseudo_dict[uid2]
+                    narration3 = sorted_narrations[i+2]
 
                 table[uid3] = {'prev2_narration': narration1,
                                'prev2_offset': round(start_times[i+2] - start_times[i],2),

llava/action/llava_inference.py

@@ -52,6 +52,9 @@ def llava_inference(
             question_type = "direct_narration"
         elif test_type == 'caption' or test_type == 'debug':
             question_type = "gpt-gt-reason"
+        elif test_type == 'temporal_cot':
+            question_type = 'temporal_cot'
+                    
         if  test_type == 'caption_then_answer':        
             caption_answer = llava_inference([video_frames], 
             tokenizer, 
@@ -73,7 +76,8 @@ def llava_inference(
                                         learn_neighbor_actions = learn_neighbor_actions,
                                         include_time_instruction= False)
 
-            question = f"You observed the video before and wrote down the notes: {caption_answer}. Now you watch the same video again and you can do better. " +  question            
+            question = f"You observed the video before and wrote down the notes: {caption_answer}. Now you watch the same video again and you can do better. " +  question                             
+            
         else:                        
             question = format_llava_prompt(DEFAULT_IMAGE_TOKEN,
                                         options,

llava/action/utils.py

@@ -235,7 +235,8 @@ def format_task_related_prompt(question, question_type, meta_data = None, perspe
         perspective_prefix = "You are seeing this video from egocentric view and you are the person. Your hands are sometimes interacting with obects. "
     elif perspective == "third_person":
         perspective_prefix = "The video is taken from egocentric view. What action is the person performing? "
-    if question_type.startswith("mc_"):
+                    
+    if question_type.startswith("mc_") or question_type == 'temporal_cot':
 
         if learn_neighbor_actions and meta_data:
             prefix = f"{perspective_prefix}\n"