import argparse import torch import os import json from tqdm import tqdm import shortuuid from ChatUniVi.constants import * from ChatUniVi.conversation import conv_templates, SeparatorStyle from ChatUniVi.model.builder import load_pretrained_model from ChatUniVi.utils import disable_torch_init from ChatUniVi.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria from PIL import Image import math from decord import VideoReader, cpu import numpy as np def read_json(file): with open(file, "r", encoding='utf-8') as f: data = json.load(f) return data def split_list(lst, n): """Split a list into n (roughly) equal-sized chunks""" chunk_size = math.ceil(len(lst) / n) # integer division return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)] def get_chunk(lst, n, k): chunks = split_list(lst, n) return chunks[k] def _get_rawvideo_dec(video_path, image_processor, max_frames=MAX_IMAGE_LENGTH, image_resolution=224, video_framerate=1, s=None, e=None): # speed up video decode via decord. video_mask = np.zeros(max_frames, dtype=np.int64) max_video_length = 0 # T x 3 x H x W video = np.zeros((max_frames, 3, image_resolution, image_resolution), dtype=np.float64) if s is None: start_time, end_time = None, None else: start_time = int(s) end_time = int(e) start_time = start_time if start_time >= 0. else 0. end_time = end_time if end_time >= 0. else 0. if start_time > end_time: start_time, end_time = end_time, start_time elif start_time == end_time: end_time = start_time + 1 if os.path.exists(video_path): vreader = VideoReader(video_path, ctx=cpu(0)) else: print(video_path) raise FileNotFoundError fps = vreader.get_avg_fps() f_start = 0 if start_time is None else int(start_time * fps) f_end = int(min(1000000000 if end_time is None else end_time * fps, len(vreader) - 1)) num_frames = f_end - f_start + 1 if num_frames > 0: # T x 3 x H x W sample_fps = int(video_framerate) t_stride = int(round(float(fps) / sample_fps)) all_pos = list(range(f_start, f_end + 1, t_stride)) if len(all_pos) > max_frames: sample_pos = [all_pos[_] for _ in np.linspace(0, len(all_pos) - 1, num=max_frames, dtype=int)] else: sample_pos = all_pos patch_images = [Image.fromarray(f) for f in vreader.get_batch(sample_pos).asnumpy()] patch_images = torch.stack([image_processor.preprocess(img, return_tensors='pt')['pixel_values'][0] for img in patch_images]) slice_len = patch_images.shape[0] max_video_length = max_video_length if max_video_length > slice_len else slice_len if slice_len < 1: pass else: video[:slice_len, ...] = patch_images return patch_images, video_mask else: print("video path: {} error.".format(video_path)) video_mask[:max_video_length] = [1] * max_video_length return torch.from_numpy(video), video_mask def eval_model(args): # Model disable_torch_init() model_path = os.path.expanduser(args.model_path) model_name = "ChatUniVi" tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, args.model_base, model_name) mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False) mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True) if mm_use_im_patch_token: tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True) if mm_use_im_start_end: tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True) model.resize_token_embeddings(len(tokenizer)) vision_tower = model.get_vision_tower() if not vision_tower.is_loaded: vision_tower.load_model() image_processor = vision_tower.image_processor if model.config.config["use_cluster"]: for n, m in model.named_modules(): m = m.to(dtype=torch.bfloat16) # Load the ground truth file with open(args.question_file) as file: gt_contents = json.load(file) answers_list = read_json(args.answers_list) answers_file = os.path.expanduser(args.answers_file) os.makedirs(os.path.dirname(answers_file), exist_ok=True) ans_file = open(answers_file, "w") video_formats = ['.mp4', '.avi', '.mov', '.mkv'] # Iterate over each sample in the ground truth file for sample in tqdm(gt_contents): sample_set = sample qs = sample['question'] # Load the video file for fmt in video_formats: # Added this line video_name = sample['video_name'] temp_path = os.path.join(args.video_folder, f"{video_name}{fmt}") if os.path.exists(temp_path): video_path = temp_path break video_name = "v_" + sample['video_name'] temp_path = os.path.join(args.video_folder, f"{video_name}{fmt}") if os.path.exists(temp_path): video_path = temp_path break # Check if the video exists if video_path is not None: # Modified this line if args.max_frames: video_frames, _ = _get_rawvideo_dec(video_path, image_processor, max_frames=args.max_frames) else: video_frames, _ = _get_rawvideo_dec(video_path, image_processor, max_frames=MAX_IMAGE_LENGTH) try: cur_prompt = qs if model.config.mm_use_im_start_end: qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN * MAX_IMAGE_LENGTH + DEFAULT_IM_END_TOKEN + '\n' + qs else: qs = DEFAULT_IMAGE_TOKEN * MAX_IMAGE_LENGTH + '\n' + qs conv = conv_templates[args.conv_mode].copy() conv.append_message(conv.roles[0], qs) conv.append_message(conv.roles[1], None) prompt = conv.get_prompt() input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze( 0).cuda() stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2 keywords = [stop_str] stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) with torch.inference_mode(): output_ids = model.generate( input_ids, images=video_frames.half().cuda(), do_sample=True, temperature=args.temperature, top_p=args.top_p, num_beams=args.num_beams, output_scores=True, return_dict_in_generate=True, max_new_tokens=1024, use_cache=True, stopping_criteria=[stopping_criteria]) output_ids = output_ids.sequences input_token_len = input_ids.shape[1] n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item() if n_diff_input_output > 0: print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids') outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0] outputs = outputs.strip() if outputs.endswith(stop_str): outputs = outputs[:-len(stop_str)] outputs = outputs.strip() ans_id = shortuuid.uuid() ans_file.write(json.dumps({"video_name": sample['video_name'], "prompt": cur_prompt, "text": outputs, "answer_id": ans_id, "model_id": model_name, "answer": sample['answer'], "metadata": {}}) + "\n") ans_file.flush() except Exception as e: print(f"Error processing video file '{video_name}': {e}") ans_file.close() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model-path", type=str, default="facebook/opt-350m") parser.add_argument("--model-base", type=str, default=None) parser.add_argument("--video-folder", type=str, default="") parser.add_argument("--question-file", type=str, default="tables/question.jsonl") parser.add_argument("--answers-list", type=str, default="tables/answers_list.json") parser.add_argument("--answers-file", type=str, default="answer.jsonl") parser.add_argument("--conv-mode", type=str, default="v1") parser.add_argument("--num-chunks", type=int, default=1) parser.add_argument("--chunk-idx", type=int, default=0) parser.add_argument("--temperature", type=float, default=0.2) parser.add_argument("--top_p", type=float, default=None) parser.add_argument("--num_beams", type=int, default=1) parser.add_argument("--max_frames", type=int, default=None) args = parser.parse_args() eval_model(args)