pku-yuangroup / video-bench Goto Github PK

import argparse
import os
import json

parser = argparse.ArgumentParser()
parser.add_argument("--dataset_name", type=str, default=None, help="The type of LLM")
parser.add_argument("--Eval_QA_root", type=str, default='./', help="folder containing QA JSON files")
parser.add_argument("--Eval_Video_root", type=str, default='./', help="folder containing video data")
parser.add_argument("--chat_conversation_output_folder", type=str, default='./Chat_results', help="")
args = parser.parse_args()

Eval_QA_root = args.Eval_QA_root
Eval_Video_root = args.Eval_Video_root
dataset_qajson = {
  "Ucfcrime": f"{Eval_QA_root}/Eval_QA/Ucfcrime_QA_new.json",
  "Youcook2": f"{Eval_QA_root}/Eval_QA/Youcook2_QA_new.json",
  "TVQA": f"{Eval_QA_root}/Eval_QA/TVQA_QA_new.json",
  "MSVD": f"{Eval_QA_root}/Eval_QA/MSVD_QA_new.json",
  "MSRVTT": f"{Eval_QA_root}/Eval_QA/MSRVTT_QA_new.json",
  "Driving-decision-making": f"{Eval_QA_root}/Eval_QA/Driving-decision-making_QA_new.json",
  "NBA": f"{Eval_QA_root}/Eval_QA/NBA_QA_new.json",
  "SQA3D": f"{Eval_QA_root}/Eval_QA/SQA3D_QA_new.json",
  "Driving-exam": f"{Eval_QA_root}/Eval_QA/Driving-exam_QA_new.json",
  "MV": f"{Eval_QA_root}/Eval_QA/MV_QA_new.json",
  "MOT": f"{Eval_QA_root}/Eval_QA/MOT_QA_new.json",
  "ActivityNet": f"{Eval_QA_root}/Eval_QA/ActivityNet_QA_new.json",
  "TGIF": f"{Eval_QA_root}/Eval_QA/TGIF_QA_new.json"
}

if args.dataset_name is None:
    dataset_name_list = list(dataset_qajson.keys())
else:
    dataset_name_list = [args.dataset_name]
    print(f'Specifically run {args.dataset_name}')
print(dataset_name_list)

os.makedirs(args.chat_conversation_output_folder, exist_ok=True)

for dataset_name in dataset_name_list:
    qa_json = dataset_qajson[dataset_name]
    print(f'Dataset name:{dataset_name}, {qa_json=}!')
    with open(qa_json, 'r', encoding='utf-8') as f:
        data = json.load(f)
        
    eval_dict = {}
    for idx, (q_id, item) in enumerate(data.items()):
        try:   
            video_id = item['video_id']
            question = item['question'] 
            if len(item['choices']) == 6:
                question += f"Choices: A.{item['choices']['A']} B.{item['choices']['B']} C.{item['choices']['C']} D.{item['choices']['D']} E.{item['choices']['E']} F.{item['choices']['F']} \n Among the six options A, B, C, D, E, F above, the one closest to the correct answer is:"
                candidates = ['A', 'B', 'C', 'D', 'E', 'F']
                candidates_long = [f" A.{item['choices']['A']}", f"B.{item['choices']['B']}", f"C.{item['choices']['C']}", f"D.{item['choices']['D']}", f"E.{item['choices']['E']}", f"F.{item['choices']['F']}"]
            elif len(item['choices']) == 5:
                question += f" A.{item['choices']['A']} B.{item['choices']['B']} C.{item['choices']['C']} D.{item['choices']['D']} E.{item['choices']['E']} \n Among the five options A, B, C, D, E above, the one closest to the correct answer is: "
                candidates = ['A', 'B', 'C', 'D', 'E']
                candidates_long = [f" A.{item['choices']['A']}", f"B.{item['choices']['B']}", f"C.{item['choices']['C']}", f"D.{item['choices']['D']}", f"E.{item['choices']['E']}"]
            elif len(item['choices']) == 4:
                question += f" A.{item['choices']['A']} B.{item['choices']['B']} C.{item['choices']['C']} D.{item['choices']['D']} \n Among the four options A, B, C, D above, the one closest to the correct answer is:"
                candidates = ['A', 'B', 'C', 'D']
                candidates_long = [f" A.{item['choices']['A']}", f"B.{item['choices']['B']}", f"C.{item['choices']['C']}", f"D.{item['choices']['D']}"]
            elif len(item['choices']) == 3:
                question += f" A.{item['choices']['A']} B.{item['choices']['B']} C.{item['choices']['C']} \n Among the three options A, B, C above, the one closest to the correct answer is: "
                candidates = ['A', 'B', 'C']
                candidates_long = [f" A.{item['choices']['A']}", f"B.{item['choices']['B']}", f"C.{item['choices']['C']}"]
            elif len(item['choices']) == 2:
                question += f" A.{item['choices']['A']} B.{item['choices']['B']} \n Among the two options A, B above, the one closest to the correct answer is: "
                candidates = ['A', 'B']
                candidates_long = [f" A.{item['choices']['A']}", f"B.{item['choices']['B']}"]
            vid_rela_path = item['vid_path']
            vid_path = os.path.join(Eval_Video_root, vid_rela_path)


            #=================================You need to change this code =========================
            # ......
            output, output_scores = ask(args, question, model, tokenizer, image_processor, vid_path)
            # ......
            #=======================================================================================

            eval_dict[q_id] = {
                'video_id': video_id,
                'question': question,
                'output_sequence': output
            }  
            print(f'q_id:{q_id}, output:{output}!\n')
        except Exception as e:
            traceback.print_exc()  
    # eval results
    eval_dataset_json = f'{args.chat_conversation_output_folder}/{dataset_name}_eval.json'
    with open(eval_dataset_json, 'w', encoding='utf-8') as f:
        json.dump(eval_dict, f, indent=2)

python Step2_chatgpt_judge.py  --model_chat_files_folder ./Chat_results  \
--apikey sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx \  # --apikey need to specify your openai apikey account
--chatgpt_judge_output_folder  ./ChatGPT_Judge

python Step3_merge_into_one_json.py  --chatgpt_judge_files_folder ./ChatGPT_Judge \
--merge_file  ./Video_Bench_Input.json