Fixed bugs

Author: yeshaokai

action/dataset.py

@@ -610,6 +610,21 @@ def get_args_parser():
 
     return parser
 
+def prepare_llava():
+
+    import warnings
+    from llava.model.builder import load_pretrained_model    
+    warnings.filterwarnings("ignore")
+    # Load the OneVision model
+    #pretrained = f"lmms-lab/llava-onevision-qwen2-{llm_size}-ov"
+    model_name = "llava_qwen"
+
+    device_map = "auto"
+    tokenizer, model, image_processor, max_length = load_pretrained_model(pretrained, None, model_name, device_map=device_map, attn_implementation="sdpa")
+
+    return tokenizer, model, image_processor, max_length
+
+
 if __name__ == '__main__':
     from moviepy.editor import ImageSequenceClip
     import torchvision
@@ -640,10 +655,24 @@ def get_args_parser():
 
     # Set up logging
     logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', filename=f'llava_ov_{args.llava_num_frames}f_{args.llm_size}.log', filemode='w')
+
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(logging.INFO)
+
+    # Set the same format for console handler as well
+    console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
+
+    # Add the console handler to the root logger
+    logging.getLogger().addHandler(console_handler)
+    
     logger = logging.getLogger(__name__)
 
+    pretrained = f"lmms-lab/llava-onevision-qwen2-{args.llm_size}-ov"
+
+    tokenizer, model, image_processor, max_length = prepare_llava()
+    
     for idx, (frames, gt) in tqdm(enumerate(val_dataloader)):
-        pred = llava_inference(frames, gt, logger, num_frames=args.llava_num_frames, llm_size=args.llm_size)
+        pred = llava_inference(frames, tokenizer, model, image_processor, max_length,  gt,  num_frames=args.llava_num_frames)
 
         # if valid letter is found in the prediction, then we will use that as the prediction
         found = False
@@ -671,4 +700,4 @@ def get_args_parser():
     # get final accuracy 
     accuracy = np.mean(gts == preds)
     logger.info(f'Final accuracy: {accuracy:.4f}')
-   
+   

action/llava_ov_inference.py

@@ -1,5 +1,5 @@
 from operator import attrgetter
-from llava.model.builder import load_pretrained_model
+
 from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
 from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IGNORE_INDEX
 from llava.conversation import conv_templates, SeparatorStyle
@@ -13,22 +13,13 @@
 import warnings
 from decord import VideoReader, cpu
 
-def llava_inference(video_frames, gt, logger, num_frames=16, llm_size='7b'):
 
-    warnings.filterwarnings("ignore")
-    # Load the OneVision model
-    pretrained = f"lmms-lab/llava-onevision-qwen2-{llm_size}-ov"
-    logger.info(f"Loading model {pretrained}")
-    model_name = "llava_qwen"
-    device = "cuda"
-    device_map = "auto"
-    tokenizer, model, image_processor, max_length = load_pretrained_model(pretrained, None, model_name, device_map=device_map, attn_implementation="sdpa")
+def llava_inference(video_frames, tokenizer, model, image_processor, max_length,  gt,  num_frames=16):
 
-    model.eval()
+    model.eval()    
+    device = "cuda"    
     video_frames = video_frames[0]
-
     temporal_stride = 16 // num_frames
-
     video_frames = video_frames[::temporal_stride]
     image_tensors = []
     frames = image_processor.preprocess(video_frames, return_tensors="pt")["pixel_values"].half().cuda()

run_EK100.sh

@@ -1,4 +1,6 @@
 python3 action/dataset.py \
     --root /media/data/haozhe/VFM/EK100/EK100_320p_15sec_30fps_libx264 \
     --train-metadata /media/data/haozhe/VFM/EK100/epic-kitchens-100-annotations/EPIC_100_train.csv \
-    --val-metadata /media/data/haozhe/VFM/EK100/epic-kitchens-100-annotations/EPIC_100_validation.csv > kitchen_test.out 2>&1
+    --val-metadata /media/data/haozhe/VFM/EK100/epic-kitchens-100-annotations/EPIC_100_validation.csv \
+    --llm_size 7b \
+    --llava_num_frames 16 > kitchen_test.out 2>&1 \