Refactor video loading function and add time instruction

Author: ZhangYuanhan-AI

playground/demo/video_demo.py

@@ -64,25 +64,31 @@ def parse_args():
     parser.add_argument("--api_key", type=str, help="OpenAI API key")
     parser.add_argument("--mm_newline_position", type=str, default="no_token")
     parser.add_argument("--force_sample", type=lambda x: (str(x).lower() == 'true'), default=False)
+    parser.add_argument("--add_time_instruction", type=str, default=False)
     return parser.parse_args()
 
 
-def load_video(video_path, args):
-    vr = VideoReader(video_path, ctx=cpu(0))
-    total_frame_num = len(vr)
-    fps = round(vr.get_avg_fps())
-    frame_idx = [i for i in range(0, len(vr), fps)]
-    # sample_fps = args.for_get_frames_num if total_frame_num > args.for_get_frames_num else total_frame_num
-    if len(frame_idx) > args.for_get_frames_num or args.force_sample:
-        sample_fps = args.for_get_frames_num
-        uniform_sampled_frames = np.linspace(0, total_frame_num - 1, sample_fps, dtype=int)
-        frame_idx = uniform_sampled_frames.tolist()
-    spare_frames = vr.get_batch(frame_idx).asnumpy()
-    # Save frames as images
-    # for i, frame in enumerate(spare_frames):
-    #     cv2.imwrite(f'{args.output_dir}/frame_{i}.jpg', cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    def load_video(video_path,args):
+        if max_frames_num == 0:
+            return np.zeros((1, 336, 336, 3))
+        vr = VideoReader(video_path, ctx=cpu(0),num_threads=1)
+        total_frame_num = len(vr)
+        video_time = total_frame_num / vr.get_avg_fps()
+        fps = round(vr.get_avg_fps()/fps)
+        frame_idx = [i for i in range(0, len(vr), fps)]
+        frame_time = [i/fps for i in frame_idx]
+        if len(frame_idx) > args.for_get_frames_num or args.force_sample:
+            sample_fps = max_frames_num
+            uniform_sampled_frames = np.linspace(0, total_frame_num - 1, sample_fps, dtype=int)
+            frame_idx = uniform_sampled_frames.tolist()
+            frame_time = [i/vr.get_avg_fps() for i in frame_idx]
+        frame_time = ",".join([f"{i:.2f}s" for i in frame_time])
+        spare_frames = vr.get_batch(frame_idx).asnumpy()
+        # import pdb;pdb.set_trace()
+
+        return spare_frames,frame_time,video_time
+
 
-    return spare_frames
 
 
 def load_video_base64(path):
@@ -177,17 +183,20 @@ def run_inference(args):
         # Check if the video exists
         if os.path.exists(video_path):
             if "gpt4v" != args.model_path:
-                video = load_video(video_path, args)
+                video,frame_time,video_time = load_video(video_path, args)
                 video = image_processor.preprocess(video, return_tensors="pt")["pixel_values"].half().cuda()
                 video = [video]
             else:
-                video = load_video_base64(video_path)
+                spare_frames,frame_time,video_time = load_video_base64(video_path)
                 interval = int(len(video) / args.for_get_frames_num)
 
         # try:
         # Run inference on the video and add the output to the list
         if "gpt4v" != args.model_path:
             qs = question
+            if args.add_time_instruction:
+                time_instruciton = f"The video lasts for {video_time:.2f} seconds, and {len(video)} frames are uniformly sampled from it. These frames are located at {frame_time}.Please answer the following questions related to this video."
+                qs = f'{time_instruciton}\n{qs}'
             if model.config.mm_use_im_start_end:
                 qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + "\n" + qs
             else: