import os
import re
import math
import json
import argparse
import warnings
import torch
import decord
import numpy as np
import transformers
from PIL import Image
from tqdm import tqdm
from decord import VideoReader, cpu
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms as T
from torchvision.transforms import functional as F
import sys
sys.path.append('./')
from videollama2.conversation import conv_templates, SeparatorStyle
from videollama2.constants import NUM_FRAMES, DEFAULT_MMODAL_TOKEN, DEFAULT_MMODAL_START_TOKEN, DEFAULT_MMODAL_END_TOKEN, MMODAL_TOKEN_INDEX
from videollama2.mm_utils import get_model_name_from_path, tokenizer_MMODAL_token, KeywordsStoppingCriteria, process_videos, expand2square
from videollama2.model.builder import load_pretrained_model
# NOTE: Ignore TypedStorage warning, which refers to this link~(https://github.com/pytorch/pytorch/issues/97207#issuecomment-1494781560)
warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
default_mm_token = DEFAULT_MMODAL_TOKEN["VIDEO"]
default_mm_start_token = DEFAULT_MMODAL_START_TOKEN["VIDEO"]
default_mm_end_token = DEFAULT_MMODAL_END_TOKEN["VIDEO"]
modal_token_index = MMODAL_TOKEN_INDEX["VIDEO"]
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]
class MVBenchDataset(Dataset):
def __init__(self, data_list, processor, num_segments=8):
self.data_list = data_list
self.decord_method = {
'video': self.read_video,
'gif': self.read_gif,
'frame': self.read_frame,
}
self.processor = processor
self.num_segments = num_segments
def __str__(self):
len_list = {}
option_list = {}
for data in self.data_list:
if data['task_type'] not in len_list:
len_list[data['task_type']] = 0
len_list[data['task_type']] += 1
if data['task_type'] not in option_list:
option_list[data['task_type']] = 0
option_list[data['task_type']] += len(data['data']['candidates'])
correct = 0
total = 0
res = f"There are {len(self.data_list)} videos as follow:\n"
for k, v in len_list.items():
correct += len_list[k]
total += option_list[k]
res += f"{v} for {k} ({option_list[k]} options => {len_list[k]/option_list[k]*100:.2f}%)\n"
correct = correct + 1 / option_list[k]
res += f"Total random accuracy: {correct/total*100:.2f}%"
return res.rstrip()
def __len__(self):
return len(self.data_list)
def get_index(self, bound, fps, max_frame, first_idx=0):
if bound:
start, end = bound[0], bound[1]
else:
start, end = -100000, 100000
start_idx = max(first_idx, round(start * fps))
end_idx = min(round(end * fps), max_frame)
seg_size = float(end_idx - start_idx) / self.num_segments
frame_indices = np.array([
int(start_idx + (seg_size / 2) + np.round(seg_size * idx))
for idx in range(self.num_segments)
])
return frame_indices
def read_video(self, video_path, bound=None):
vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
max_frame = len(vr) - 1
fps = float(vr.get_avg_fps())
images_group = list()
frame_indices = self.get_index(bound, fps, max_frame, first_idx=0)
for frame_index in frame_indices:
img = Image.fromarray(vr[frame_index].asnumpy())
images_group.append(img)
# images_group = [expand2square(img, tuple(int(x*255) for x in self.processor.image_mean)) for img in images_group]
torch_imgs = self.processor(images_group, return_tensors='pt')['pixel_values']
return torch_imgs
def read_gif(self, video_path, bound=None, fps=25):
gif = imageio.get_reader(video_path)
max_frame = len(gif) - 1
images_group = list()
frame_indices = self.get_index(bound, fps, max_frame, first_idx=0)
for index, frame in enumerate(gif):
if index in frame_indices:
img = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)
img = Image.fromarray(img)
images_group.append(img)
# images_group = [expand2square(img, tuple(int(x*255) for x in self.processor.image_mean)) for img in images_group]
torch_imgs = self.processor(images_group, return_tensors='pt')['pixel_values']
return torch_imgs
def read_frame(self, video_path, bound=None, fps=3):
max_frame = len(os.listdir(video_path))
images_group = list()
frame_indices = self.get_index(bound, fps, max_frame, first_idx=1) # frame_idx starts from 1
for frame_index in frame_indices:
img = Image.open(os.path.join(video_path, f"{frame_index:05d}.jpg"))
images_group.append(img)
# images_group = [expand2square(img, tuple(int(x*255) for x in self.processor.image_mean)) for img in images_group]
torch_imgs = self.processor.preprocess(images_group, return_tensors='pt')['pixel_values']
return torch_imgs
def qa_template(self, data):
question = f"Question: {data['question']}\n"
question += "Options:\n"
answer = data['answer']
answer_idx = -1
for idx, c in enumerate(data['candidates']):
question += f"({chr(ord('A') + idx)}) {c}\n"
if c == answer:
answer_idx = idx
question = question.rstrip()
answer = f"({chr(ord('A') + answer_idx)}) {answer}"
return question, answer
def __getitem__(self, idx):
decord_method = self.decord_method[self.data_list[idx]['data_type']]
bound = None
if self.data_list[idx]['bound']:
bound = (
self.data_list[idx]['data']['start'],
self.data_list[idx]['data']['end'],
)
video_path = os.path.join(self.data_list[idx]['prefix'], self.data_list[idx]['data']['video'])
torch_imgs = decord_method(video_path, bound)
question = self.data_list[idx]['data']['question']
options = self.data_list[idx]['data']['candidates']
answer = self.data_list[idx]['data']['answer']
task_type = self.data_list[idx]['task_type']
# question, answer = self.qa_template(self.data_list[idx]['data'])
answer_idx = -1
letters = []
options_string = ''
for option_idx, c in enumerate(options):
letters.append(f"{chr(ord('A') + option_idx)}")
options_string += f"({chr(ord('A') + option_idx)}) {c}\n"
if c == answer:
answer_idx = option_idx
option_question = f'Question: {question}\nOptions:\n{options_string}Answer with the option\'s letter from the given choices directly and only give the best option.'
return {
'video': torch_imgs,
'video_path': video_path,
'question': option_question,
'letters': ','.join(letters),
'answer_idx': answer_idx,
'task_type': task_type
}
tasks = {
"Action Sequence": ("action_sequence.json", "star/Charades_v1_480/", "video", True), # has start & end
"Action Prediction": ("action_prediction.json", "star/Charades_v1_480/", "video", True), # has start & end
"Action Antonym": ("action_antonym.json", "ssv2_video/", "video", False),
"Fine-grained Action": ("fine_grained_action.json", "Moments_in_Time_Raw/videos/", "video", False),
"Unexpected Action": ("unexpected_action.json", "FunQA_test/test/", "video", False),
"Object Existence": ("object_existence.json", "clevrer/video_validation/", "video", False),
"Object Interaction": ("object_interaction.json", "star/Charades_v1_480/", "video", True), # has start & end
"Object Shuffle": ("object_shuffle.json", "perception/videos/", "video", False),
"Moving Direction": ("moving_direction.json", "clevrer/video_validation/", "video", False),
"Action Localization": ("action_localization.json", "sta/sta_video/", "video", True), # has start & end
"Scene Transition": ("scene_transition.json", "scene_qa/video/", "video", False),
"Action Count": ("action_count.json", "perception/videos/", "video", False),
"Moving Count": ("moving_count.json", "clevrer/video_validation/", "video", False),
"Moving Attribute": ("moving_attribute.json", "clevrer/video_validation/", "video", False),
"State Change": ("state_change.json", "perception/videos/", "video", False),
"Fine-grained Pose": ("fine_grained_pose.json", "nturgbd/", "video", False),
"Character Order": ("character_order.json", "perception/videos/", "video", False),
"Egocentric Navigation": ("egocentric_navigation.json", "vlnqa/", "video", False),
"Episodic Reasoning": ("episodic_reasoning.json", "tvqa/frames_fps3_hq/", "frame", True), # has start & end, read frame
"Counterfactual Inference": ("counterfactual_inference.json", "clevrer/video_validation/", "video", False),
}
def build_mvbench_eval(args, processor, num_frames):
data_list = []
for task_name, task in tasks.items():
json_file = os.path.join(args.question_file, task[0])
vis_folder = os.path.join(args.video_folder, task[1])
with open(json_file, 'r') as f:
json_data = json.load(f)
for data in json_data:
data_list.append({
'task_type': task_name,
'prefix': vis_folder,
'data_type': task[2],
'bound': task[3],
'data': data
})
data_list = get_chunk(data_list, args.num_chunks, args.chunk_idx)
dataset = MVBenchDataset(data_list, processor, num_segments=num_frames)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
return dataloader
def mvbench_dump(ans_file, line, outputs):
for idx, output in enumerate(outputs):
vid = line['video_path'][idx]
task_type = line['task_type'][idx]
letters = line['letters'][idx].split(',')
answer_idx = line['answer_idx'][idx].item()
pred_answer = re.findall(f'[\(,\ ]*[{letters[0]}-{letters[-1]}][\),\ ]*', output)
if len(pred_answer) == 0:
pred_idx = (answer_idx + 1) % len(letters)
else:
pred_answer = pred_answer[0].strip()
if pred_answer.startswith('('):
pred_answer = pred_answer.strip('()')
pred_idx = letters.index(pred_answer)
ans_file.write(json.dumps({"vid": vid, "task_type": task_type, "pred": pred_idx, "gt": answer_idx}) + '\n')
class NextoeDataset(Dataset):
video_formats = ['.mp4', '.avi', '.mov', '.mkv']
def __init__(self, data_list, processor, num_segments=8):
self.data_list = data_list
self.processor = processor
self.num_segments = num_segments
def __len__(self):
return len(self.data_list)
def __getitem__(self, idx):
line = self.data_list[idx]
video_name = line['video']
question = line['question']
answer = line['answer']
for fmt in self.video_formats: # Added this line
temp_path = os.path.join(args.video_folder, f"{video_name}{fmt}")
if os.path.exists(temp_path):
video_path = temp_path
break
decord_vr = VideoReader(uri=video_path, ctx=cpu(0))
frames = decord_vr.get_batch(np.linspace(0, len(decord_vr) - 1, 8, dtype=int)).asnumpy()
video_tensor = self.processor.preprocess(frames, return_tensors='pt')['pixel_values'] # do not pad for video frames
wrapped_question = f'Question: {question}\nAnswer the question using a single word or a short phrase with multiple words.'
return {
'video': video_tensor,
'question': wrapped_question,
'answer': answer,
'qid': line['qid']
}
def build_nextoe_eval(args, processor, num_frames):
questions = json.load(open(args.question_file, "r"))
questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
dataset = NextoeDataset(questions, processor, num_segments=num_frames)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
return dataloader
def nextoe_dump(ans_file, line, outputs):
for idx, output in enumerate(outputs):
vid, qid = line['qid'][idx].split('_')
ans_file.write(json.dumps({"vid": vid, "qid": qid, "prediction": output}) + '\n')
class NextqaDataset(Dataset):
video_formats = ['.mp4', '.avi', '.mov', '.mkv']
def __init__(self, data_list, processor, num_segments=8):
self.data_list = data_list
self.processor = processor
self.num_segments = num_segments
def __len__(self):
return len(self.data_list)
def __getitem__(self, idx):
line = self.data_list[idx]
video_name = line['video']
question = line['question']
answer = line['answer']
for fmt in self.video_formats: # Added this line
temp_path = os.path.join(args.video_folder, f"{video_name}{fmt}")
if os.path.exists(temp_path):
video_path = temp_path
break
decord_vr = VideoReader(uri=video_path, ctx=cpu(0))
frames = decord_vr.get_batch(np.linspace(0, len(decord_vr) - 1, 8, dtype=int)).asnumpy()
video_tensor = self.processor.preprocess(frames, return_tensors='pt')['pixel_values'] # do not pad for video frames
assert line['num_option'] == 5
a0 = line['a0']
a1 = line['a1']
a2 = line['a2']
a3 = line['a3']
a4 = line['a4']
option_question = f'Question: {question}\nOptions:\n(A) {a0}\n(B) {a1}\n(C) {a2}\n(D) {a3}\n(E) {a4}\nAnswer with the option\'s letter from the given choices directly and only give the best option.'
return {
'video': video_tensor,
'question': option_question,
'answer': answer,
'qid': line['qid']
}
def build_nextqa_eval(args, processor, num_frames):
questions = json.load(open(args.question_file, "r"))
questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
dataset = NextqaDataset(questions, processor, num_segments=num_frames)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
return dataloader
def nextqa_dump(ans_file, line, outputs):
for idx, output in enumerate(outputs):
qid = line['qid'][idx]
answer = line['answer'][idx].item()
letters = ['A', 'B', 'C', 'D', 'E']
pred_answer = re.findall('[\(,\ ]*[A-E][\),\ ]*', output)
if len(pred_answer) == 0:
pred_idx = 2
else:
pred_answer = pred_answer[0].strip()
if pred_answer.startswith('('):
pred_answer = pred_answer.strip('()')
pred_idx = letters.index(pred_answer)
ans_file.write(json.dumps({"id": qid, "prediction": pred_idx, "answer": answer}) + '\n')
class EgoschemaDataset(Dataset):
video_formats = ['.mp4', '.avi', '.mov', '.mkv']
def __init__(self, data_list, processor, num_segments=8):
self.data_list = data_list
self.processor = processor
self.num_segments = num_segments
def __len__(self):
return len(self.data_list)
def __getitem__(self, idx):
line = self.data_list[idx]
q_uid = line['q_uid']
for fmt in self.video_formats: # Added this line
temp_path = os.path.join(args.video_folder, f"{q_uid}{fmt}")
if os.path.exists(temp_path):
video_path = temp_path
break
decord_vr = VideoReader(uri=video_path, ctx=cpu(0))
frames = decord_vr.get_batch(np.linspace(0, len(decord_vr) - 1, self.num_segments, dtype=int)).asnumpy()
video_tensor = self.processor.preprocess(frames, return_tensors='pt')['pixel_values'] # do not pad for video frames
question = line['question']
a0 = line['option 0']
a1 = line['option 1']
a2 = line['option 2']
a3 = line['option 3']
a4 = line['option 4']
axs = [a0, a1, a2, a3, a4]
ops = ['(A)', '(B)', '(C)', '(D)', '(E)']
option_question = f'Question: {question}\nOptions:\n(A) {a0}\n(B) {a1}\n(C) {a2}\n(D) {a3}\n(E) {a4}\n.Answer with the option\'s letter from the given choices directly and only give the best option.'
return {
'q_uid': q_uid,
'video': video_tensor,
'question': option_question,
}
def build_egoschema_eval(args, processor, num_frames):
questions = json.load(open(args.question_file, "r"))
questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
dataset = EgoschemaDataset(questions, processor, num_segments=num_frames)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
return dataloader
def egoschema_dump(ans_file, line, outputs):
for idx, output in enumerate(outputs):
q_uid = line['q_uid'][idx]
letters = ['A', 'B', 'C', 'D', 'E']
pred_answer = re.findall('[\(\ ]*[A-E][\)\ ]*', output)
if len(pred_answer) == 0:
pred_idx = 2
else:
pred_answer = pred_answer[0].strip()
# if pred_answer.startswith('('):
pred_answer = pred_answer.strip('()')
pred_idx = letters.index(pred_answer)
ans_file.write(f'{q_uid}, {pred_idx}\n')
def get_model_output(model, video_tensor, tokenizer, questions, conv_mode="v1", device='cuda'):
input_ids = []
modal_list = []
for qs in questions:
if model.config.mm_use_im_start_end:
qs = default_mm_start_token + default_mm_token + default_mm_end_token + "\n" + qs
else:
qs = default_mm_token + "\n" + qs
conv = conv_templates[conv_mode].copy()
conv.append_message(conv.roles[0], qs)
conv.append_message(conv.roles[1], None)
prompt = conv.get_prompt()
input_id = tokenizer_MMODAL_token(prompt, tokenizer, modal_token_index, return_tensors='pt')
input_ids.append(input_id)
modal_list.append("video")
# left pad sequence
input_ids = torch.nn.utils.rnn.pad_sequence(
[x.flip(dims=[0]) for x in input_ids],
batch_first=True,
padding_value=tokenizer.pad_token_id).flip(dims=[1]).to(device)
attention_mask=input_ids.ne(tokenizer.pad_token_id).to(device)
video_tensor = video_tensor.half().to(args.device)
with torch.inference_mode():
output_ids = model.generate(
input_ids,
attention_mask=attention_mask,
images_or_videos=video_tensor,
modal_list=modal_list,
do_sample=False,
max_new_tokens=1024,
use_cache=True,
pad_token_id=tokenizer.eos_token_id)
outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
return outputs
def run_inference(args):
"""
Run inference on ActivityNet QA DataSet using the Video-ChatGPT model.
Args:
args: Command-line arguments.
"""
# Initialize the model
model_name = get_model_name_from_path(args.model_path)
tokenizer, model, processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name)
num_frames = model.config.num_frames if hasattr(model.config, "num_frames") else NUM_FRAMES
answer_file = os.path.expanduser(args.answer_file)
os.makedirs(os.path.dirname(answer_file), exist_ok=True)
ans_file = open(answer_file, "w")
output_list = [] # List to store the output results
if args.dataset == 'mvbench':
val_loader = build_mvbench_eval(args, processor, num_frames)
elif args.dataset == 'nextoe':
val_loader = build_nextoe_eval(args, processor, num_frames)
elif args.dataset == 'nextqa':
val_loader = build_nextqa_eval(args, processor, num_frames)
elif args.dataset == 'egoschema':
val_loader = build_egoschema_eval(args, processor, num_frames)
else:
raise NotImplementedError(f"Dataset {args.dataset} not implemented.")
# Iterate over each sample in the ground truth file
for i, line in enumerate(tqdm(val_loader)):
video_tensor = line['video']
questions = line['question']
outputs = get_model_output(model, video_tensor, tokenizer, questions, args.conv_mode, args.device)
if args.dataset == 'mvbench':
mvbench_dump(ans_file, line, outputs)
elif args.dataset == 'nextoe':
nextoe_dump(ans_file, line, outputs)
elif args.dataset == 'nextqa':
nextqa_dump(ans_file, line, outputs)
elif args.dataset == 'egoschema':
egoschema_dump(ans_file, line, outputs)
else:
raise NotImplementedError(f"Dataset {args.dataset} not implemented.")
ans_file.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Multiple-Choice Video QA Evaluation Script.')
parser.add_argument('--dataset', help='Dataset to evaluate on.', required=True)
parser.add_argument('--model-path', help='', required=True)
parser.add_argument('--model_base', help='', default=None, type=str, required=False)
parser.add_argument('--video-folder', help='Directory containing video files.', required=True)
parser.add_argument('--question-file', help='Path to the ground truth file containing question.', required=True)
parser.add_argument('--answer-file', help='Path to the ground truth file containing answers.', required=True)
parser.add_argument("--conv-mode", type=str, default="llava_v1")
parser.add_argument("--num-chunks", type=int, default=1)
parser.add_argument("--chunk-idx", type=int, default=0)
parser.add_argument("--device", type=str, required=False, default='cuda:0')
parser.add_argument("--model_max_length", type=int, required=False, default=2048)
parser.add_argument("--batch-size", type=int, default=1)
parser.add_argument("--num-workers", type=int, default=8)
args = parser.parse_args()
run_inference(args)