Initial upload of all project files

.ruff_cache/.gitignore

@@ -0,0 +1,2 @@
+# Automatically created by ruff.
+*

.ruff_cache/CACHEDIR.TAG

@@ -0,0 +1 @@
+Signature: 8a477f597d28d172789f06886806bc55

compute_video_emb.py

@@ -0,0 +1,154 @@
+import os
+import torch
+import argparse
+from PIL import Image
+from tqdm import tqdm
+from transformers import AutoModel, AutoProcessor
+import torch.multiprocessing as mp
+from torch.utils.data import Dataset, DataLoader
+import glob
+
+# --- 配置 ---
+MODEL_ID = "/mnt/bn/ziyang-storage-cloudnative-hl/huggingface/siglip-so400m-patch14-384"
+BATCH_SIZE = 1024  # 根据你的 GPU VRAM 调整
+
+def parse_arguments():
+    """解析命令行参数"""
+    parser = argparse.ArgumentParser(
+        description="步骤 1: 使用 SigLIP (多GPU) 预计算所有视频帧的嵌入."
+    )
+    parser.add_argument(
+        "--frames-path",
+        "-fp",
+        type=str,
+        required=True,
+        help="包含所有视频帧文件夹的基础目录的绝对路径。",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        required=True,
+        help="用于保存嵌入.pt文件的输出目录路径。",
+    )
+    return parser.parse_args()
+
+class FrameDataset(Dataset):
+    """一个用于高效加载视频帧的PyTorch Dataset"""
+    def __init__(self, frame_paths):
+        self.frame_paths = frame_paths
+
+    def __len__(self):
+        return len(self.frame_paths)
+
+    def __getitem__(self, idx):
+        path = self.frame_paths[idx]
+        try:
+            image = Image.open(path).convert("RGB")
+            return image
+        except Exception:
+            return None
+
+def collate_fn(batch):
+    """自定义collate函数，用于从批次中过滤掉None值"""
+    batch = [item for item in batch if item is not None]
+    if not batch:
+        return None
+    return batch
+
+def process_video_chunk(args_tuple):
+    """
+    工作函数，用于在特定GPU上处理一批视频。
+    """
+    video_dirs_chunk, frames_base_path, gpu_id, output_dir = args_tuple
+    device = f"cuda:{gpu_id}"
+
+    # 在工作进程中为指定的GPU加载模型和处理器
+    model = AutoModel.from_pretrained(MODEL_ID).to(device).eval()
+    processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=True)
+
+    progress_bar = tqdm(video_dirs_chunk, position=gpu_id, desc=f"GPU-{gpu_id}")
+
+    for video_dir in progress_bar:
+        video_name = os.path.basename(video_dir)
+        output_path = os.path.join(output_dir, f"{video_name}.pt")
+
+        # 如果文件已存在，则跳过，以支持断点续算
+        if os.path.exists(output_path):
+            progress_bar.write(f"Skipping {video_name}, embeddings already exist.")
+            continue
+        
+        frame_files = [f for f in os.listdir(video_dir) if f.endswith(".jpg")]
+        if not frame_files:
+            continue
+        frame_files.sort(key=lambda x: int(x.split("_")[1].split(".")[0]))
+        frame_paths = [os.path.join(video_dir, f) for f in frame_files]
+
+        try:
+            with torch.no_grad():
+                dataset = FrameDataset(frame_paths)
+                loader = DataLoader(
+                    dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0,
+                    pin_memory=True, collate_fn=collate_fn
+                )
+
+                all_frame_embeddings = []
+                for image_batch in loader:
+                    if image_batch is None:
+                        continue
+                    
+                    image_inputs = processor(images=image_batch, return_tensors="pt").to(device)
+                    frame_embeddings = model.get_image_features(**image_inputs)
+                    all_frame_embeddings.append(frame_embeddings)
+
+                if not all_frame_embeddings:
+                    continue
+
+                all_frame_embeddings = torch.cat(all_frame_embeddings, dim=0)
+
+                # 将张量移动到CPU以便保存，避免后续加载时出现CUDA问题
+                data_to_save = {
+                    'filenames': frame_files,
+                    'embeddings': all_frame_embeddings.cpu()
+                }
+                torch.save(data_to_save, output_path)
+
+        except Exception as e:
+            progress_bar.write(f"Error on GPU-{gpu_id} for video '{video_name}': {e}")
+
+def main():
+    """主函数，用于协调多GPU处理"""
+    args = parse_arguments()
+
+    num_gpus = torch.cuda.device_count()
+    if num_gpus == 0:
+        print("错误: 未找到启用CUDA的GPU。正在退出。")
+        exit(1)
+    
+    print(f"找到 {num_gpus} 个GPU。开始并行处理...")
+
+    # 创建输出目录
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # 获取所有唯一的视频目录
+    video_dirs = [d for d in glob.glob(os.path.join(args.frames_path, '*')) if os.path.isdir(d)]
+    
+    if not video_dirs:
+        print(f"错误: 在 {args.frames_path} 中未找到视频目录。")
+        return
+
+    # 将视频目录分成块，每个GPU一块
+    chunk_size = (len(video_dirs) + num_gpus - 1) // num_gpus
+    video_chunks = [video_dirs[i:i + chunk_size] for i in range(0, len(video_dirs), chunk_size)]
+    
+    # 为每个工作进程准���参数
+    process_args = [(video_chunks[i], args.frames_path, i, args.output_dir) for i in range(len(video_chunks))]
+
+    with mp.Pool(processes=num_gpus) as pool:
+        pool.map(process_video_chunk, process_args)
+
+    print("\n所有视频帧嵌入已计算并保存。")
+
+if __name__ == "__main__":
+    mp.set_start_method('spawn', force=True)
+    main()

main_adaptive_sampling.py

@@ -0,0 +1,447 @@
+import os
+import json
+import base64
+import argparse
+import time
+import re
+from datetime import datetime
+from functools import partial
+from openai import AzureOpenAI, OpenAI
+from volcenginesdkarkruntime import Ark
+from multiprocessing import Pool, Manager, Lock
+
+# New prompt template for multiple-choice questions with reasoning
+REASONING_MULTIPLE_CHOICE_TEMPLATE = """
+You are an AI assistant evaluating video frames to answer a multiple-choice question.
+The user will provide you with a set of video frames and a question with several options (e.g., A, B, C, D).
+
+First, provide a step-by-step reasoning process that analyzes the video frames and leads to your conclusion.
+After your reasoning, provide the final answer in a JSON block. The JSON object must contain a single key "answer" with the value being one of 'A', 'B', 'C', or 'D'.
+
+Your output should follow this format exactly:
+<Your step-by-step reasoning here>
+```json
+{"answer": "A"}
+```
+Do not include any other text after the JSON block.
+"""
+
+
+def parse_arguments():
+    """
+    Parse command line arguments for evaluation configuration.
+
+    Returns:
+        argparse.Namespace: Parsed command line arguments
+    """
+    parser = argparse.ArgumentParser(
+        description="Video QA Evaluation with Pre-computed Similarity Frame Selection"
+    )
+
+    # Model configuration
+    parser.add_argument(
+        "--target-model",
+        "-tm",
+        type=str,
+        required=True,
+        help="Model to be evaluated (e.g., gpt-4o, gpt-4-vision-preview)",
+    )
+
+    # Data configuration
+    parser.add_argument(
+        "--frame-num",
+        "-fn",
+        type=int,
+        default=32,
+        help="Number of most similar frames to select for each video (default: 32)",
+    )
+    parser.add_argument(
+        "--frames-path",
+        "-fp",
+        type=str,
+        required=True,
+        help="Absolute path to the base directory containing video frame folders.",
+    )
+    parser.add_argument(
+        "--data-file",
+        "-df",
+        type=str,
+        required=True,
+        help="Absolute path to the JSON file containing the evaluation dataset.",
+    )
+    # --- MODIFIED ARGUMENT ---
+    parser.add_argument(
+        "--similarity-file",
+        "-sf",
+        type=str,
+        required=True,
+        help="Absolute path to the pre-computed similarity JSON file (e.g., lv_bench_similarity.json).",
+    )
+
+    # Processing configuration
+    parser.add_argument(
+        "--max-retry-times",
+        "-mr",
+        type=int,
+        default=10,
+        help="Maximum number of retries for API calls (default: 10)",
+    )
+    parser.add_argument(
+        "--pool-processes",
+        "-pp",
+        type=int,
+        default=20,
+        help="Number of parallel processes for evaluation (default: 20)",
+    )
+
+    # API configuration
+    parser.add_argument(
+        "--base_url", type=str, required=True, help="Azure OpenAI endpoint URL."
+    )
+    parser.add_argument(
+        "--api_key", type=str, required=True, help="Azure OpenAI API key."
+    )
+
+    return parser.parse_args()
+
+
+def save_json_file(data, output_file):
+    """
+    Save data to a JSON file.
+    """
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4)
+
+
+def extract_json_from_response(response):
+    """
+    Extracts a JSON object from a string that contains reasoning followed by a tagged JSON block.
+    """
+    if not response:
+        return None
+    try:
+        match = re.search(r"```json\s*(\{.*?\})\s*```", response, re.DOTALL)
+        if match:
+            json_str = match.group(1)
+            return json.loads(json_str)
+        return None
+    except (json.JSONDecodeError, IndexError):
+        return None
+
+
+def calculate_metrics(results):
+    """
+    Calculate evaluation metrics from the results.
+    """
+    total_samples = len(results)
+    if total_samples == 0:
+        return {
+            "total_samples": 0,
+            "answered_samples": 0,
+            "correct_answers": 0,
+            "accuracy": 0.0,
+        }
+
+    answered_samples = sum(1 for x in results if x.get("model_answer") is not None)
+    correct_answers = sum(1 for x in results if x.get("is_correct"))
+
+    accuracy = correct_answers / answered_samples if answered_samples > 0 else 0.0
+
+    return {
+        "total_samples": total_samples,
+        "answered_samples": answered_samples,
+        "correct_answers": correct_answers,
+        "accuracy": accuracy,
+    }
+
+
+def call_single_model(client, messages, model, item_id, max_retry_times):
+    """
+    Make a single API call to the specified model with retry logic.
+    """
+    if "doubao" in model:
+        max_tokens = 32768
+    else:
+        max_tokens = 65535
+    retry_times = 0
+    while retry_times < max_retry_times:
+        try:
+            completion = client.chat.completions.create(
+                model=model, messages=messages, max_tokens=max_tokens
+            )
+            return completion.choices[0].message.content
+        except Exception as e:
+            retry_times += 1
+            print(
+                f"Error processing item {item_id} with model {model}: {str(e)}. Retrying ({retry_times}/{max_retry_times})..."
+            )
+            if retry_times == max_retry_times:
+                error_log_file = f"error_log_{model.replace('/', '_')}.txt"
+                with open(error_log_file, "a") as f:
+                    f.write(
+                        f"Error processing item {item_id} with model {model} after {max_retry_times} retries: {str(e)}\n"
+                    )
+                return None
+            time.sleep(5)
+
+
+def evaluate_single_item(
+    data_item, frames, target_model, api_key, base_url, max_retry_times
+):
+    """
+    Evaluate a single data item using the target model.
+    """
+    if "ark" in base_url:
+        client = Ark(base_url=base_url, api_key=api_key)
+    elif "aliyun" in base_url or "127.0.0.1" in base_url:
+        client = OpenAI(api_key=api_key, base_url=base_url)
+    else:
+        client = AzureOpenAI(
+            api_version="2023-05-15", api_key=api_key, azure_endpoint=base_url
+        )
+
+    messages = [
+        {"role": "system", "content": REASONING_MULTIPLE_CHOICE_TEMPLATE},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Here are the video frames:"},
+                *frames,
+                {"type": "text", "text": f"Question: {data_item['question']}"},
+            ],
+        },
+    ]
+
+    response = call_single_model(
+        client, messages, target_model, data_item["key"], max_retry_times
+    )
+
+    is_correct = False
+    model_answer_cleaned = None
+    parsed_json = None
+
+    if response:
+        parsed_json = extract_json_from_response(response)
+        if parsed_json and "answer" in parsed_json:
+            model_answer_cleaned = str(parsed_json["answer"]).strip().upper()
+            gold_answer = data_item["answer"].strip().upper()
+            if model_answer_cleaned == gold_answer:
+                is_correct = True
+
+    return {
+        **data_item,
+        "model_reasoning_and_answer": response,
+        "model_answer_raw": parsed_json.get("answer") if parsed_json else None,
+        "model_answer": model_answer_cleaned,
+        "is_correct": is_correct,
+    }
+
+
+def encode_image(image_path):
+    """
+    Encode an image file to base64 string.
+    """
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+# --- MODIFIED: New function for selecting frames based on pre-computed similarity file ---
+def process_frames_from_similarity_file(
+    frames_base_path, frame_num, data_item, similarity_data
+):
+    """
+    Select and encode the top N frames using a pre-computed similarity file.
+    """
+    item_key = data_item["key"]
+    question_uid = str(data_item["uid"])
+
+    # Retrieve the sorted list of frame filenames for the current question
+    sorted_filenames = similarity_data.get(question_uid)
+
+    if not sorted_filenames:
+        print(
+            f"Warning: No similarity data found for question UID '{question_uid}', skipping."
+        )
+        return []
+
+    try:
+        # Select the top N filenames
+        num_frames_to_select = min(frame_num, len(sorted_filenames))
+        selected_filenames = sorted_filenames[:num_frames_to_select]
+        selected_ids = [int(f.split(".")[0].split("_")[-1]) for f in selected_filenames]
+        selected_ids = sorted(selected_ids)
+        selected_filenames = [f"frame_{i:06d}.jpg" for i in selected_ids]
+
+        # Construct full paths for the selected frames
+        video_frames_path = os.path.join(frames_base_path, item_key)
+        sampled_paths = [os.path.join(video_frames_path, f) for f in selected_filenames]
+
+        # Encode the selected frames
+        base64_images = [encode_image(path) for path in sampled_paths]
+
+        return [
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/jpeg;base64,{b64_img}"},
+            }
+            for b64_img in base64_images
+        ]
+    except Exception as e:
+        print(f"Error during frame processing for key '{item_key}': {e}")
+        return []
+
+
+def process_single_data(
+    data_item,
+    args,
+    shared_results,
+    progress_counter,
+    total_items,
+    locks,
+    similarity_data,
+):
+    """
+    Process a single data item in a multiprocessing context.
+    """
+    item_key = data_item["key"]
+    try:
+        # --- MODIFIED: Call the new frame selection function ---
+        frames = process_frames_from_similarity_file(
+            args.frames_path, args.frame_num, data_item, similarity_data
+        )
+
+        if not frames:
+            raise ValueError(
+                f"No frames were processed from similarity file for key '{item_key}'"
+            )
+
+        result = evaluate_single_item(
+            data_item,
+            frames,
+            args.target_model,
+            args.api_key,
+            args.base_url,
+            args.max_retry_times,
+        )
+
+        if result is not None:
+            with locks["results"]:
+                shared_results.append(result)
+                data_filename_base = os.path.splitext(os.path.basename(args.data_file))[
+                    0
+                ]
+                model_name_safe = args.target_model.replace("/", "_")
+                output_prefix = f"{model_name_safe}_{data_filename_base}_{args.frame_num}frames_precomputed_similar"
+                results_output_file = f"{output_prefix}_results.json"
+                save_json_file(list(shared_results), results_output_file)
+
+    except Exception as e:
+        print(f"Error processing video key {item_key}: {str(e)}")
+        with locks["file"]:
+            error_log_file = f"error_log_{args.target_model.replace('/', '_')}.txt"
+            with open(error_log_file, "a") as f:
+                f.write(f"Critical error processing video key {item_key}: {str(e)}\n")
+    finally:
+        with locks["counter"]:
+            progress_counter.value += 1
+            print(
+                f"\rProcessed: {progress_counter.value}/{total_items} videos...",
+                end="",
+                flush=True,
+            )
+
+
+def load_test_data(json_file):
+    """
+    Load test data from a JSON file.
+    """
+    try:
+        with open(json_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        print(f"Error: Data file not found at {json_file}")
+        exit(1)
+    except json.JSONDecodeError:
+        print(f"Error: Could not decode JSON from {json_file}")
+        exit(1)
+
+
+def main():
+    """
+    Main function to run the video QA evaluation framework.
+    """
+    args = parse_arguments()
+
+    print("--- Evaluation Configuration ---")
+    print(f"Target Model: {args.target_model}")
+    print(f"Frames to Sample (by pre-computed similarity): {args.frame_num}")
+    print(f"Frames Base Path: {args.frames_path}")
+    print(f"Similarity File: {args.similarity_file}")  # Print new arg
+    print(f"Data File: {args.data_file}")
+    print(f"Parallel Processes: {args.pool_processes}")
+    print("---------------------------------")
+
+    error_log_file = f"error_log_{args.target_model.replace('/', '_')}.txt"
+    with open(error_log_file, "w") as f:
+        f.write(
+            f"=== Error Log Started at {datetime.now()} for model {args.target_model} ===\n"
+        )
+
+    data_filename_base = os.path.splitext(os.path.basename(args.data_file))[0]
+    model_name_safe = args.target_model.replace("/", "_")
+    output_prefix = f"{model_name_safe}_{data_filename_base}_{args.frame_num}frames_precomputed_similar"
+
+    results_output_file = f"{output_prefix}_results.json"
+    metrics_output_file = f"{output_prefix}_metrics.json"
+
+    # Load test data and similarity data
+    test_data = load_test_data(args.data_file)
+    try:
+        with open(args.similarity_file, "r", encoding="utf-8") as f:
+            similarity_data = json.load(f)
+    except FileNotFoundError:
+        print(f"Error: Similarity file not found at {args.similarity_file}")
+        exit(1)
+
+    total_videos = len(test_data)
+    print(f"\nLoaded {total_videos} videos to process.")
+
+    with Manager() as manager:
+        shared_results = manager.list()
+        progress_counter = manager.Value("i", 0)
+        locks = {
+            "results": manager.Lock(),
+            "file": manager.Lock(),
+            "counter": manager.Lock(),
+        }
+
+        # Create a partial function with fixed arguments for the worker pool
+        process_func = partial(
+            process_single_data,
+            args=args,
+            shared_results=shared_results,
+            progress_counter=progress_counter,
+            total_items=total_videos,
+            locks=locks,
+            similarity_data=similarity_data,
+        )
+
+        # Run processing in parallel
+        with Pool(processes=args.pool_processes) as pool:
+            pool.map(process_func, test_data)
+
+        all_results = list(shared_results)
+
+    print(f"\n\nProcessing complete for model: {args.target_model}")
+
+    final_metrics = calculate_metrics(all_results)
+    save_json_file(final_metrics, metrics_output_file)
+    print(f"\nMetrics saved to: {metrics_output_file}")
+    print(json.dumps(final_metrics, indent=4))
+
+    save_json_file(all_results, results_output_file)
+    print(f"Detailed results saved to: {results_output_file}")
+
+
+if __name__ == "__main__":
+    main()

main_agent.py

@@ -0,0 +1,521 @@
+import os
+import json
+import base64
+import argparse
+import time
+import re
+import traceback
+from datetime import datetime
+from functools import partial
+from openai import AzureOpenAI, OpenAI
+from volcenginesdkarkruntime import Ark
+import concurrent.futures
+from tqdm import tqdm
+
+# New system prompt for the agent
+AGENT_SYSTEM_PROMPT = """
+You are an intelligent AI assistant specialized in video question answering.
+Your task is to answer a multiple-choice question based on a video.
+
+You must use the `get_frames_by_id` tool to request specific frames to view.
+You will be told the total number of frames available in the video (e.g., "The video has 1250 frames, numbered 1 to 1250.").
+
+Your strategy should be efficient:
+1.  Based on the task query, think about which part of the video will be related, and then get the frames of this part. If the query’s description is fairly general and you can’t effectively infer the temporal regions where the target visual evidence might appear, you can first uniformly sample some frames for analysis to identify the time intervals where the target visual evidence is likely to appear.
+2.  Analyze the retrieved frames and the user's question.
+3.  If you don't have enough information, form a hypothesis about where the answer might be and use the tool again to request more specific frames from that segment.
+4.  Continue this process of reasoning and tool use until you are confident in your answer. Avoid requesting all frames at once.
+5.  Please make sure that you find the relevant visual cues and then answer the question instead of guessing the answer.
+6.  You can access 10 frames at most in each tool call.
+
+Please note that if you have insufficient visual information at the beginning, you can first sample more frames uniformly to understand the video (e.g., sampling 10 frames per tool call). You can then gradually refine the subsequent steps and adopt a coarse-to-fine strategy overall.
+For example, the question is "What is the main subject of the video?"
+You can first sample 10 frames uniformly from the video (e.g., frame 100, 200, ..., 1200).
+After analyzing these frames, you might notice that the main subject is a person in the middle of the screen (between frame 500 and 600).
+You can then sample more frames from this region (e.g., frame 500, 520, ..., 590) to get more detailed information.
+Finally, you can reason based on the visual cues you have gathered and provide the final answer.
+This process might be multi-turn.
+
+After your reasoning, provide the final answer in a JSON block. The JSON object must contain a single key "answer" with the value being one of 'A', 'B', 'C', or 'D'.
+
+Remember that You can access 10 frames at most in each tool call.
+
+Your output should follow this format exactly:
+<Your step-by-step reasoning here>
+```json
+{"answer": "X"}
+```
+Do not include any other text after the JSON block.
+"""
+
+# Tool schema for the get_frames_by_id function
+GET_FRAMES_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "get_frames_by_id",
+        "description": "Retrieves specific video frames by their numerical IDs. Use this to get visual information from the video.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "frame_ids": {
+                    "type": "array",
+                    "items": {"type": "integer"},
+                    "description": "A list of frame numbers to retrieve. You can access 10 frames at most in each tool call.",
+                },
+            },
+            "required": ["frame_ids"],
+        },
+    },
+}
+
+
+def parse_arguments():
+    """
+    Parse command line arguments for evaluation configuration.
+    """
+    parser = argparse.ArgumentParser(
+        description="Video QA Evaluation Framework with Agentic Frame Selection (Refactored)"
+    )
+    parser.add_argument(
+        "--target-model",
+        "-tm",
+        type=str,
+        required=True,
+        help="Model to be evaluated (e.g., gpt-4o)",
+    )
+    parser.add_argument(
+        "--frames-path",
+        "-fp",
+        type=str,
+        required=True,
+        help="Absolute path to the base directory for video frames.",
+    )
+    parser.add_argument(
+        "--data-file",
+        "-df",
+        type=str,
+        required=True,
+        help="Absolute path to the JSON evaluation dataset.",
+    )
+    parser.add_argument(
+        "--max-retry-times",
+        "-mr",
+        type=int,
+        default=10,
+        help="Maximum retries for API calls.",
+    )
+    parser.add_argument(
+        "--pool-processes",
+        "-pp",
+        type=int,
+        default=20,
+        help="Number of parallel processes.",
+    )
+    parser.add_argument(
+        "--base_url", type=str, required=True, help="Azure OpenAI endpoint URL."
+    )
+    parser.add_argument(
+        "--api_key", type=str, required=True, help="Azure OpenAI API key."
+    )
+    return parser.parse_args()
+
+
+def save_json_file(data, output_file):
+    """Saves data to a JSON file."""
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4)
+
+
+def extract_json_from_response(response):
+    """Extracts a JSON object from a model's response string."""
+    if not response:
+        return None
+    match = re.search(r"```json\s*(\{.*?\})\s*```", response, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(1))
+        except (json.JSONDecodeError, IndexError):
+            return None
+    return None
+
+
+def calculate_metrics(results):
+    """Calculates accuracy and other metrics from evaluation results."""
+    # Filter out potential error results before calculating
+    valid_results = [r for r in results if "error" not in r]
+    total_samples = len(valid_results)
+
+    if total_samples == 0:
+        return {
+            "total_samples": 0,
+            "answered_samples": 0,
+            "correct_answers": 0,
+            "accuracy": 0.0,
+        }
+
+    answered_samples = sum(
+        1 for x in valid_results if x.get("model_answer") is not None
+    )
+    correct_answers = sum(1 for x in valid_results if x.get("is_correct"))
+    accuracy = correct_answers / answered_samples if answered_samples > 0 else 0.0
+
+    return {
+        "total_samples": total_samples,
+        "answered_samples": answered_samples,
+        "correct_answers": correct_answers,
+        "accuracy": accuracy,
+    }
+
+
+def call_single_model(client, messages, model, item_id, max_retry_times, tools=None):
+    """Makes a single API call with retry logic and tool support."""
+    if "o4" in model:
+        params = {"model": model, "messages": messages, "max_tokens": 65535}
+    elif "Qwen" in model:
+        params = {
+            "model": model,
+            "messages": messages,
+            "max_tokens": 2048,
+            "temperature": 0,
+        }
+    else:
+        params = {"model": model, "messages": messages, "max_tokens": 32768}
+    if tools:
+        params["tools"] = tools
+        params["tool_choice"] = "auto"
+
+    retry_times = 0
+    while retry_times < max_retry_times:
+        try:
+            completion = client.chat.completions.create(**params)
+            return completion.choices[0].message
+        except Exception as e:
+            retry_times += 1
+            print(
+                f"API Error for item {item_id}: {str(e)}. Retrying ({retry_times}/{max_retry_times})..."
+            )
+            if retry_times == max_retry_times:
+                # Instead of writing to a file here, we'll let the worker return the error
+                raise e  # Reraise the exception to be caught by the worker's main try-except block
+            time.sleep(5)
+
+
+def get_frames_by_id(frame_ids: list, all_frame_paths: list):
+    """Tool implementation: Retrieves and formats frames based on a list of IDs."""
+    retrieved_frames = []
+    frame_map = {
+        int(re.search(r"frame_(\d+)\.jpg", os.path.basename(p)).group(1)): p
+        for p in all_frame_paths
+        if re.search(r"frame_(\d+)\.jpg", os.path.basename(p))
+    }
+
+    for fid in frame_ids:
+        path = frame_map.get(fid)
+        if path and os.path.exists(path):
+            b64_image = encode_image(path)
+            retrieved_frames.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/jpeg;base64,{b64_image}"},
+                }
+            )
+    return retrieved_frames
+
+
+def evaluate_single_item_agentic(
+    data_item, all_frame_paths, target_model, api_key, base_url, max_retry_times
+):
+    """Evaluates a single item using an agentic loop for dynamic frame selection."""
+    if "ark" in base_url:
+        client = Ark(
+            base_url=base_url,
+            api_key=api_key,
+        )
+    elif "aliyun" in base_url or "127.0.0.1" in base_url:
+        client = OpenAI(api_key=api_key, base_url=base_url)
+    else:
+        client = AzureOpenAI(
+            api_version="2023-05-15", api_key=api_key, azure_endpoint=base_url
+        )
+
+    tools = [GET_FRAMES_TOOL_SCHEMA]
+    available_functions = {"get_frames_by_id": get_frames_by_id}
+
+    total_frames = len(all_frame_paths)
+    minutes = data_item["video_info"]["duration_minutes"]
+    seconds = int(minutes * 60)
+    initial_prompt = (
+        f"The video has {total_frames} frames, numbered 1 to {total_frames}. This video is {seconds} seconds long. "
+        f"Please answer the following question:\n{data_item['question']}"
+    )
+
+    messages = [
+        {"role": "system", "content": AGENT_SYSTEM_PROMPT},
+        {"role": "user", "content": initial_prompt},
+    ]
+
+    response_content = None
+    max_tool_calls = 10
+
+    for i in range(max_tool_calls):
+        response_message = call_single_model(
+            client,
+            messages,
+            target_model,
+            data_item["key"],
+            max_retry_times,
+            tools=tools,
+        )
+        if response_message is None:
+            return None
+
+        messages.append(response_message.model_dump())
+
+        if response_message.tool_calls:
+            for tool_call in response_message.tool_calls:
+                function_name = tool_call.function.name
+                function_to_call = available_functions.get(function_name)
+                if function_to_call:
+                    function_args = json.loads(tool_call.function.arguments)
+                    retrieved_frames = function_to_call(
+                        **function_args, all_frame_paths=all_frame_paths
+                    )
+                    tool_response_content = [
+                        {
+                            "type": "text",
+                            "text": f"Here are the frames you requested (IDs: {function_args.get('frame_ids', [])}).",
+                        }
+                    ]
+                    tool_response_content.extend(retrieved_frames)
+                    messages.append(
+                        {
+                            "tool_call_id": tool_call.id,
+                            "role": "tool",
+                            "name": function_name,
+                            "content": json.dumps(
+                                {
+                                    "status": "success",
+                                    "retrieved_frame_count": len(retrieved_frames),
+                                }
+                            ),
+                        }
+                    )
+                    messages.append({"role": "user", "content": tool_response_content})
+        else:
+            response_content = response_message.content
+            break
+
+    if response_content is None and response_message and response_message.tool_calls:
+        print(
+            f"\nMax tool calls reached for item {data_item['key']}. Forcing a final answer."
+        )
+        final_prompt = "You have reached the maximum number of tool calls. Please provide a final answer in the specified JSON format based on the information you have gathered so far."
+        messages.append({"role": "user", "content": final_prompt})
+        final_response_message = call_single_model(
+            client,
+            messages,
+            target_model,
+            data_item["key"],
+            max_retry_times,
+            tools=None,
+        )
+        if final_response_message:
+            messages.append(final_response_message)
+            response_content = final_response_message.content
+    elif response_content is None and response_message:
+        response_content = response_message.content
+
+    is_correct = False
+    model_answer_cleaned = None
+    parsed_json = extract_json_from_response(response_content)
+    if parsed_json and "answer" in parsed_json:
+        model_answer_cleaned = str(parsed_json["answer"]).strip().upper()
+        gold_answer = data_item["answer"].strip().upper()
+        if model_answer_cleaned == gold_answer:
+            is_correct = True
+    return {
+        **data_item,
+        "agent_conversation": [
+            msg.model_dump() if hasattr(msg, "model_dump") else msg for msg in messages
+        ],
+        "model_reasoning_and_answer": response_content,
+        "model_answer": model_answer_cleaned,
+        "is_correct": is_correct,
+    }
+
+
+def encode_image(image_path):
+    """Encodes an image file to a base64 string."""
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+def process_single_data(data_item, args):
+    """
+    Main processing function for a single video.
+    This function is executed by each worker process. It is self-contained.
+    """
+    item_key = data_item["key"]
+    try:
+        specific_frames_path = os.path.join(args.frames_path, item_key)
+        if not os.path.isdir(specific_frames_path):
+            raise FileNotFoundError(f"Frame directory not found for key '{item_key}'")
+
+        all_frame_paths = sorted(
+            [
+                os.path.join(specific_frames_path, f)
+                for f in os.listdir(specific_frames_path)
+                if f.endswith(".jpg")
+            ],
+            key=lambda x: int(re.search(r"frame_(\d+)\.jpg", x).group(1)),
+        )
+
+        if not all_frame_paths:
+            raise FileNotFoundError(f"No frames found for key '{item_key}'")
+
+        # The core evaluation logic is called here
+        result = evaluate_single_item_agentic(
+            data_item,
+            all_frame_paths,
+            args.target_model,
+            args.api_key,
+            args.base_url,
+            args.max_retry_times,
+        )
+        return result
+
+    except Exception as e:
+        # If any error occurs, catch it and return an error dictionary.
+        # This prevents the worker process from crashing and allows the main
+        # process to log the error gracefully.
+        print(f"\nCRITICAL ERROR on key {item_key}: {str(e)}")
+        traceback.print_exc()
+        return {
+            "key": item_key,
+            "uid": data_item.get("uid"),
+            "error": str(e),
+            "traceback": traceback.format_exc(),
+        }
+
+
+def load_test_data(json_file):
+    """Loads the evaluation data from a JSON file."""
+    try:
+        with open(json_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        print(f"Error: Data file not found: {json_file}")
+        exit(1)
+    except json.JSONDecodeError:
+        print(f"Error: Malformed JSON in {json_file}")
+        exit(1)
+
+
+def main():
+    """Main function to orchestrate the evaluation framework."""
+    args = parse_arguments()
+
+    print("--- Agentic Video QA Evaluation (Refactored) ---")
+    print(f"Target Model: {args.target_model}")
+    print(f"Frames Base Path: {args.frames_path}")
+    print(f"Data File: {args.data_file}")
+
+    model_name_safe = args.target_model.replace("/", "_")
+    data_filename_base = os.path.splitext(os.path.basename(args.data_file))[0]
+
+    output_prefix = f"{model_name_safe}_{data_filename_base}_agent_results"
+    results_output_file = f"{output_prefix}.json"
+    metrics_output_file = f"{output_prefix}_metrics.json"
+    error_log_file = f"{output_prefix}_errors.log"
+
+    with open(error_log_file, "a", encoding="utf-8") as f:
+        f.write(
+            f"\n=== Log Session Started at {datetime.now()} for {args.target_model} ===\n"
+        )
+
+    all_test_data = load_test_data(args.data_file)
+    completed_ids = set()
+    existing_results = []
+
+    if os.path.exists(results_output_file):
+        try:
+            with open(results_output_file, "r", encoding="utf-8") as f:
+                existing_results = json.load(f)
+            if isinstance(existing_results, list):
+                completed_ids = {
+                    item["uid"] for item in existing_results if "uid" in item
+                }
+                print(
+                    f"Found {len(completed_ids)} completed tasks in '{results_output_file}'. Resuming..."
+                )
+            else:
+                existing_results = []
+        except (json.JSONDecodeError, IOError) as e:
+            print(f"Warning: Could not read results file: {e}. Starting fresh.")
+            existing_results = []
+
+    tasks_to_process = [
+        item for item in all_test_data if item.get("uid") not in completed_ids
+    ]
+
+    if not tasks_to_process:
+        print("All tasks are already completed. Calculating final metrics.")
+        final_metrics = calculate_metrics(existing_results)
+        save_json_file(final_metrics, metrics_output_file)
+        print(f"\nFinal metrics saved to: {metrics_output_file}")
+        print(json.dumps(final_metrics, indent=4))
+        return
+
+    print(
+        f"Total tasks: {len(all_test_data)}. Completed: {len(completed_ids)}. To process: {len(tasks_to_process)}."
+    )
+
+    # This list will hold all results, both old and new.
+    all_results = list(existing_results)
+
+    # Using ProcessPoolExecutor for robust, modern multiprocessing.
+    with concurrent.futures.ProcessPoolExecutor(
+        max_workers=args.pool_processes
+    ) as executor:
+        # partial is used to pass the constant `args` to each call of process_single_data
+        func = partial(process_single_data, args=args)
+
+        # executor.map processes the tasks in parallel.
+        # tqdm provides a progress bar.
+        results_iterator = executor.map(func, tasks_to_process)
+
+        for result in tqdm(
+            results_iterator, total=len(tasks_to_process), desc="Processing Videos"
+        ):
+            if result:
+                if "error" in result:
+                    # Log errors centrally
+                    with open(error_log_file, "a", encoding="utf-8") as f:
+                        f.write(f"Error on key {result.get('key', 'N/A')}:\n")
+                        f.write(f"  Error: {result['error']}\n")
+                        f.write(f"  Traceback: {result['traceback']}\n---\n")
+
+                # Append every result (success or error) to the main list
+                all_results.append(result)
+
+                # Periodically save results for resilience
+                if len(all_results) % 10 == 0:
+                    save_json_file(all_results, results_output_file)
+
+    print("\n\nProcessing complete.")
+
+    # Final save of all combined results
+    save_json_file(all_results, results_output_file)
+    print(f"Detailed results saved to: {results_output_file}")
+
+    # Calculate and save final metrics
+    final_metrics = calculate_metrics(all_results)
+    save_json_file(final_metrics, metrics_output_file)
+    print(f"\nMetrics saved to: {metrics_output_file}")
+    print(json.dumps(final_metrics, indent=4))
+
+
+if __name__ == "__main__":
+    # To run this script, you'll need to install tqdm:
+    # pip install tqdm
+    main()

main_i2i_ret.py

@@ -0,0 +1,548 @@
+import os
+import json
+import base64
+import argparse
+import time
+import re
+import traceback
+import uuid
+import multiprocessing
+import concurrent.futures
+from datetime import datetime
+from functools import partial
+
+import requests
+import torch
+from PIL import Image
+from tqdm import tqdm
+from openai import AzureOpenAI, OpenAI
+from volcenginesdkarkruntime import Ark
+from transformers import AutoModel, AutoProcessor
+from torch.nn.functional import cosine_similarity
+
+# --- Model and Configuration Constants ---
+
+# SigLIP model for generating image embeddings
+SIGLIP_MODEL_ID = "/mnt/bn/ziyang-storage-cloudnative-hl/huggingface/siglip-so400m-patch14-384"
+# Number of Top-K frames to retrieve for each generated image
+TOP_K_FRAMES = 8
+
+# --- Prompt Templates ---
+
+# Step 1: System prompt for VLM to analyze video and question, then generate image creation requests
+# The goal of this prompt is not to answer the question, but to plan which keyframes need to be "seen"
+STEP_1_PLANNING_PROMPT = """
+You are a professional video analyst. Your task is to analyze a question and a few initial video sample frames, then plan what keyframes you need to see to answer the question.
+
+Do not answer the question directly. Your output must be a JSON array, where each object represents a keyframe you wish to generate.
+Each object must contain the following two keys:
+1. `reference_image_id`: An integer representing the ID of a frame already provided to you that you wish to use as a generation reference. This ID must be one of the IDs provided by the user.
+2. `prompt`: A detailed text description to tell the image generation model what kind of scene to draw.
+
+For example, if the question is "Where did the man in the red shirt eventually go?", you might generate the following JSON:
+```json
+[
+  {
+    "reference_image_id": 120,
+    "prompt": "A man in a red shirt is walking towards an open door, with a background similar to the reference image."
+  },
+  {
+    "reference_image_id": 120,
+    "prompt": "A man in a red shirt has already walked out the door, and the door is closing, with a background similar to the reference image."
+  }
+]
+```
+Your output must strictly adhere to this JSON format.
+"""
+
+# Step 3: System prompt for VLM to perform final reasoning and answer based on all retrieved keyframes
+STEP_3_FINAL_ANSWER_PROMPT = """
+You are an AI video question-answering assistant.
+The user will provide you with a series of keyframes retrieved from a video and a question.
+
+First, please provide a step-by-step reasoning process, analyzing these keyframes and deriving your conclusion.
+After your reasoning, provide the final answer. The answer must be in a JSON code block, and the JSON object must contain a key "answer" with a value of one of 'A', 'B', 'C', or 'D'.
+
+Your output format must be strictly as follows:
+<Your step-by-step reasoning process>
+```json
+{"answer": "A"}
+```
+Do not include any other text after the JSON block.
+"""
+
+
+def parse_arguments():
+    """Parse command-line arguments"""
+    parser = argparse.ArgumentParser(
+        description="Image Retrieval-based Video QA Workflow"
+    )
+    # Model Configuration
+    parser.add_argument(
+        "--target-model", "-tm", type=str, required=True, help="VLM model for inference (e.g., gpt-4o)"
+    )
+    # Data Path Configuration
+    parser.add_argument(
+        "--frames-path", "-fp", type=str, required=True, help="Root directory containing video frame folders"
+    )
+    parser.add_argument(
+        "--data-file", "-df", type=str, required=True, help="JSON data file containing evaluation questions"
+    )
+    parser.add_argument(
+        "--embeddings-path", "-ep", type=str, required=True, help="Directory containing pre-computed embeddings for all video frames"
+    )
+    parser.add_argument(
+        "--output-path", "-op", type=str, default="./results_image_retrieval", help="Directory to store all outputs and generated images"
+    )
+    # Workflow Parameters
+    parser.add_argument(
+        "--initial-frames-num", "-ifn", type=int, default=8, help="Number of initial uniformly sampled frames for Step 1"
+    )
+    # Execution Configuration
+    parser.add_argument(
+        "--max-retry-times", "-mr", type=int, default=10, help="Maximum number of retries for API calls"
+    )
+    parser.add_argument(
+        "--pool-processes", "-pp", type=int, default=10, help="Number of parallel processes"
+    )
+    # API Credentials
+    parser.add_argument(
+        "--base_url", type=str, required=True, help="API Endpoint URL for the VLM model"
+    )
+    parser.add_argument(
+        "--api_key", type=str, required=True, help="API Key for the VLM model"
+    )
+    return parser.parse_args()
+
+
+def save_json_file(data, output_file):
+    """Save data to a JSON file"""
+    os.makedirs(os.path.dirname(output_file), exist_ok=True)
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4, ensure_ascii=False)
+
+
+def extract_json_from_response(response, is_list=False):
+    """Extract a JSON object or list from the model's response text"""
+    if not response:
+        return None
+    # The regular expression supports both JSON objects `{...}` and lists `[...]`
+    pattern = r"```json\s*([\{\[].*?[\]\}])\s*```"
+    match = re.search(pattern, response, re.DOTALL)
+    if match:
+        json_str = match.group(1)
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError:
+            print(f"JSON parsing failed: {json_str}")
+            return None
+    return None
+
+
+def calculate_metrics(results):
+    """Calculate accuracy and other metrics from evaluation results"""
+    valid_results = [r for r in results if "error" not in r]
+    total_samples = len(valid_results)
+    if total_samples == 0: return {"accuracy": 0.0}
+    
+    answered = sum(1 for x in valid_results if x.get("model_answer") is not None)
+    correct = sum(1 for x in valid_results if x.get("is_correct"))
+    accuracy = correct / answered if answered > 0 else 0.0
+    
+    return {
+        "total_samples": total_samples,
+        "answered_samples": answered,
+        "correct_answers": correct,
+        "accuracy": accuracy,
+    }
+
+
+def call_vlm_api(client, messages, model, item_id, max_retry_times, json_schema=None):
+    """Call VLM API, with support for retries and structured output"""
+    params = {"model": model, "messages": messages, "max_tokens": 4096}
+    if json_schema:
+        params["response_format"] = {"type": "json_object", "schema": json_schema}
+
+    for retry in range(max_retry_times):
+        try:
+            completion = client.chat.completions.create(**params)
+            return completion.choices[0].message.content
+        except Exception as e:
+            print(f"API Error (item {item_id}): {e}. Retrying ({retry + 1}/{max_retry_times})...")
+            if retry == max_retry_times - 1:
+                raise e
+            time.sleep(5)
+
+
+def generate_image(reference_image_id, prompt, all_frame_paths, output_dir, generation_idx):
+    """Call the image generation API to create a new frame"""
+    print(f"\n[Image Generation] Using Prompt: '{prompt}'")
+    ark_api_key = os.environ.get("ARK_API_KEY")
+    if not ark_api_key:
+        raise ValueError("Environment variable ARK_API_KEY is not set.")
+
+    client = Ark(base_url="https://ark.cn-beijing.volces.com/api/v3", api_key=ark_api_key)
+    
+    ref_image_path = all_frame_paths.get(reference_image_id)
+    if not ref_image_path or not os.path.exists(ref_image_path):
+        raise FileNotFoundError(f"Reference image ID {reference_image_id} not found.")
+
+    try:
+        ref_image_b64 = encode_image(ref_image_path)
+        ref_image_data_uri = f"data:image/jpeg;base64,{ref_image_b64}"
+        
+        response = client.images.generate(
+            model="doubao-seedream-4-0-250828",
+            prompt=prompt,
+            image=ref_image_data_uri,
+            size="1024x1024",
+            response_format="url",
+            watermark=False,
+        )
+        image_url = response.data[0].url
+        
+        image_content = requests.get(image_url, timeout=60).content
+        
+        new_frame_filename = f"generated_frame_{generation_idx}_ref_{reference_image_id}.jpg"
+        new_frame_path = os.path.join(output_dir, new_frame_filename)
+        
+        with open(new_frame_path, "wb") as f:
+            f.write(image_content)
+            
+        print(f"[Image Generation Success] Image saved to: {new_frame_path}")
+        return new_frame_path
+    except Exception as e:
+        print(f"Image generation or download failed: {e}")
+        traceback.print_exc()
+        return None
+
+def retrieve_frames_by_image_embedding(
+    image_path, video_embeddings_data, request_queue, results_dict, k
+):
+    """Retrieve Top-K similar frames from the video using an image embedding"""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    frame_filenames = video_embeddings_data["filenames"]
+    frame_embeddings = video_embeddings_data["embeddings"].to(device)
+
+    # 1. Send request to the embedding server process
+    request_id = str(uuid.uuid4())
+    request_queue.put((request_id, image_path))
+
+    # 2. Wait for the result
+    while request_id not in results_dict:
+        time.sleep(0.05)
+    query_embedding = results_dict.pop(request_id).to(device)
+
+    # 3. Calculate similarity and find Top-K frames
+    with torch.no_grad():
+        similarities = cosine_similarity(query_embedding, frame_embeddings)
+        top_k_indices = torch.topk(similarities, k=min(k, len(frame_filenames)), dim=-1).indices.cpu()
+        
+        # Extract absolute paths for the frames from the filenames
+        video_frame_dir = os.path.dirname(frame_filenames[0])
+        top_k_paths = [os.path.join(video_frame_dir, video_embeddings_data['filenames'][i]) for i in top_k_indices]
+
+    return top_k_paths
+
+def embedding_server_process(model_id, device, request_queue, results_dict):
+    """
+    An independent server process that loads the SigLIP model and handles image embedding requests from worker processes.
+    """
+    print(f"Embedding server started (PID: {os.getpid()})...")
+    model = AutoModel.from_pretrained(model_id).to(device).eval()
+    processor = AutoProcessor.from_pretrained(model_id)
+    print("SigLIP model loaded in the embedding server.")
+
+    while True:
+        try:
+            request_id, image_path = request_queue.get()
+            if image_path == "STOP":
+                print("Embedding server received stop signal, shutting down.")
+                break
+            
+            with torch.no_grad():
+                image = Image.open(image_path).convert("RGB")
+                inputs = processor(images=[image], return_tensors="pt").to(device)
+                image_features = model.get_image_features(**inputs)
+                results_dict[request_id] = image_features.cpu()
+
+        except Exception as e:
+            print(f"Error in embedding server: {e}")
+            traceback.print_exc()
+
+
+def encode_image(image_path):
+    """Encode an image file to a Base64 string"""
+    with open(image_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def uniformly_sample_frames_and_encode(frames_dir, num_frames):
+    """Uniformly sample frames and encode them, while also returning a mapping of frame IDs to paths"""
+    if not os.path.isdir(frames_dir): return [], {}
+
+    frame_files = sorted(
+        [f for f in os.listdir(frames_dir) if f.endswith(".jpg")],
+        key=lambda x: int(re.search(r"frame_(\d+)\.jpg", x).group(1)),
+    )
+    if not frame_files: return [], {}
+
+    indices = [int(i * len(frame_files) / num_frames) for i in range(num_frames)]
+    sampled_files = [frame_files[i] for i in indices]
+
+    frame_path_map, encoded_frames = {}, []
+    for f in sampled_files:
+        path = os.path.join(frames_dir, f)
+        frame_id = int(re.search(r"frame_(\d+)\.jpg", f).group(1))
+        
+        encoded_frames.extend([
+            {"type": "text", "text": f"This is Frame ID: {frame_id}"},
+            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(path)}"}}
+        ])
+        frame_path_map[frame_id] = path
+    return encoded_frames, frame_path_map
+
+
+def run_workflow_for_item(
+    data_item, args, request_queue, results_dict
+):
+    """Execute the complete three-step workflow for a single data item"""
+    item_key = data_item["key"]
+    print(f"\n--- Starting processing for video: {item_key} ---")
+    
+    # Create a separate output directory for each video's generated images
+    generated_images_dir = os.path.join(args.output_path, "generated_images", item_key)
+    os.makedirs(generated_images_dir, exist_ok=True)
+    
+    # Initialize VLM client
+    if "ark" in args.base_url:
+        client = Ark(base_url=args.base_url, api_key=args.api_key)
+    elif "aliyun" in args.base_url or "127.0.0.1" in args.base_url:
+        client = OpenAI(api_key=args.api_key, base_url=args.base_url)
+    else:
+        client = AzureOpenAI(api_version="2023-05-15", api_key=args.api_key, azure_endpoint=args.base_url)
+    
+    # --- Step 1: Initial understanding and generating "keyframe profile" requests ---
+    print(f"[{item_key}] Step 1: Uniformly sampling and generating keyframe creation requests...")
+    video_frames_path = os.path.join(args.frames_path, item_key)
+    initial_frames_encoded, initial_frame_paths = uniformly_sample_frames_and_encode(
+        video_frames_path, args.initial_frames_num
+    )
+    if not initial_frames_encoded:
+        raise FileNotFoundError(f"Initial frames not found for video {item_key}.")
+
+    planning_messages = [
+        {"role": "system", "content": STEP_1_PLANNING_PROMPT},
+        {"role": "user", "content": [
+            {"type": "text", "text": "Here are the initial sample frames and the question:"},
+            *initial_frames_encoded,
+            {"type": "text", "text": f"Question: {data_item['question']}"}
+        ]}
+    ]
+    
+    # Define JSON Schema for structured output
+    planning_schema = {
+        "type": "array",
+        "items": {
+            "type": "object",
+            "properties": {
+                "reference_image_id": {"type": "integer"},
+                "prompt": {"type": "string"}
+            },
+            "required": ["reference_image_id", "prompt"]
+        }
+    }
+    
+    raw_planning_response = call_vlm_api(client, planning_messages, args.target_model, item_key, args.max_retry_times)
+    image_generation_requests = extract_json_from_response(raw_planning_response, is_list=True)
+    
+    if not image_generation_requests or not isinstance(image_generation_requests, list):
+        raise ValueError(f"Step 1 failed to generate valid JSON-formatted image generation requests. Response: {raw_planning_response}")
+    
+    print(f"[{item_key}] Successfully generated {len(image_generation_requests)} keyframe generation requests.")
+
+    # --- Validate and correct reference image IDs ---
+    valid_ids = list(initial_frame_paths.keys())
+    if not valid_ids:
+        raise ValueError(f"No valid initial frame IDs found for video {item_key}.")
+        
+    for req in image_generation_requests:
+        original_id = req.get("reference_image_id")
+        if original_id not in valid_ids:
+            closest_id = min(valid_ids, key=lambda valid_id: abs(valid_id - original_id))
+            print(f"Warning: Model generated a non-existent reference_image_id: {original_id}. Substituting with the closest valid ID: {closest_id}.")
+            req["reference_image_id"] = closest_id
+
+    # --- Step 2: Generate images and perform similarity retrieval ---
+    print(f"[{item_key}] Step 2: Generating images and retrieving similar frames...")
+    all_retrieved_frame_paths = set()
+    generated_image_paths = []
+    video_embedding_file = os.path.join(args.embeddings_path, f"{item_key}.pt")
+    if not os.path.exists(video_embedding_file):
+        raise FileNotFoundError(f"Embedding file for video {item_key} not found: {video_embedding_file}")
+    video_embeddings_data = torch.load(video_embedding_file, map_location="cpu")
+    
+    # Correct path issue, ensure filenames in the embedding file are absolute paths
+    video_frame_dir_for_embeddings = os.path.join(args.frames_path, item_key)
+    video_embeddings_data['filenames'] = [os.path.join(video_frame_dir_for_embeddings, os.path.basename(f)) for f in video_embeddings_data['filenames']]
+
+
+    for i, req in enumerate(image_generation_requests):
+        # 2a. Generate image
+        generated_path = generate_image(
+            reference_image_id=req["reference_image_id"],
+            prompt=req["prompt"],
+            all_frame_paths=initial_frame_paths,
+            output_dir=generated_images_dir,
+            generation_idx=i + 1,
+        )
+
+        path_for_retrieval = None
+        if generated_path:
+            generated_image_paths.append(generated_path)
+            path_for_retrieval = generated_path
+        else:
+            print(f"Warning: Generation failed for image {i+1}. Using its reference image (ID: {req['reference_image_id']}) for retrieval instead.")
+            path_for_retrieval = initial_frame_paths.get(req["reference_image_id"])
+
+        if not path_for_retrieval:
+            print(f"Error: Could not find a path for retrieval for request {i+1}. Skipping.")
+            continue
+        
+        # 2b. Retrieve frames via image embedding
+        retrieved_paths = retrieve_frames_by_image_embedding(
+            path_for_retrieval, video_embeddings_data, request_queue, results_dict, k=TOP_K_FRAMES
+        )
+        all_retrieved_frame_paths.update(retrieved_paths)
+        print(f"[{item_key}] Retrieval {i+1}/{len(image_generation_requests)} complete, found {len(retrieved_paths)} frames.")
+
+    if not all_retrieved_frame_paths:
+        raise ValueError(f"Failed to retrieve any frames for video {item_key}.")
+        
+    print(f"[{item_key}] Step 2 complete. Retrieved a total of {len(all_retrieved_frame_paths)} unique keyframes.")
+    
+    # --- Step 3: Consolidate keyframes for final reasoning ---
+    print(f"[{item_key}] Step 3: Consolidating keyframes for final reasoning...")
+    final_frames_encoded = []
+    for path in sorted(list(all_retrieved_frame_paths)):
+        final_frames_encoded.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(path)}"}})
+
+    final_messages = [
+        {"role": "system", "content": STEP_3_FINAL_ANSWER_PROMPT},
+        {"role": "user", "content": [
+            {"type": "text", "text": "Here are all the keyframes retrieved for you. Please answer the question based on them."},
+            *final_frames_encoded,
+            {"type": "text", "text": f"Question: {data_item['question']}"}
+        ]}
+    ]
+
+    final_response_text = call_vlm_api(client, final_messages, args.target_model, item_key, args.max_retry_times)
+    
+    # --- Consolidating Results ---
+    parsed_answer = extract_json_from_response(final_response_text)
+    model_answer = parsed_answer.get("answer", "").strip().upper() if parsed_answer else None
+    is_correct = (model_answer == data_item["answer"].strip().upper()) if model_answer else False
+    
+    result = {
+        **data_item,
+        "workflow_steps": {
+            "step1_planning_requests": image_generation_requests,
+            "step2_generated_images": generated_image_paths,
+            "step2_retrieved_frame_paths": sorted(list(all_retrieved_frame_paths)),
+            "step3_final_reasoning_and_answer": final_response_text,
+        },
+        "model_answer": model_answer,
+        "is_correct": is_correct,
+    }
+    return result
+
+
+def process_single_data_wrapper(data_item, args, request_queue, results_dict):
+    """Wrapper function to process a single data item, used for exception handling"""
+    try:
+        return run_workflow_for_item(data_item, args, request_queue, results_dict)
+    except Exception as e:
+        print(f"\nA critical error occurred while processing video {data_item['key']}: {e}")
+        traceback.print_exc()
+        return {
+            "key": data_item['key'],
+            "uid": data_item.get('uid'),
+            "error": str(e),
+            "traceback": traceback.format_exc(),
+        }
+
+def main():
+    """Main function to orchestrate the entire evaluation workflow"""
+    args = parse_arguments()
+    print("--- Image Retrieval-based Video QA Workflow Starting ---")
+    print(f"Evaluating Model: {args.target_model}, Dataset: {args.data_file}")
+    
+    try:
+        multiprocessing.set_start_method("spawn", force=True)
+    except RuntimeError:
+        pass # Start method already set
+
+    os.makedirs(args.output_path, exist_ok=True)
+    
+    # Define output file paths
+    model_safe_name = args.target_model.replace("/", "_")
+    data_filename_base = os.path.splitext(os.path.basename(args.data_file))[0]
+    output_prefix = f"{model_safe_name}_{data_filename_base}_image_retrieval_{args.initial_frames_num}frames"
+    
+    results_file = os.path.join(args.output_path, f"{output_prefix}_results.json")
+    metrics_file = os.path.join(args.output_path, f"{output_prefix}_metrics.json")
+    
+    test_data = load_test_data(args.data_file)
+    all_results = []
+    
+    with multiprocessing.Manager() as manager:
+        request_queue = manager.Queue()
+        results_dict = manager.dict()
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        embedding_server = multiprocessing.Process(
+            target=embedding_server_process,
+            args=(SIGLIP_MODEL_ID, device, request_queue, results_dict),
+        )
+        embedding_server.start()
+        
+        # Wait for the embedding server model to load
+        time.sleep(15)
+
+        with concurrent.futures.ProcessPoolExecutor(max_workers=args.pool_processes) as executor:
+            func = partial(
+                process_single_data_wrapper,
+                args=args,
+                request_queue=request_queue,
+                results_dict=results_dict
+            )
+            
+            results_iterator = executor.map(func, test_data)
+            
+            for result in tqdm(results_iterator, total=len(test_data), desc="Processing Videos"):
+                if result:
+                    all_results.append(result)
+                    # Save results every 10 videos to prevent data loss from interruptions
+                    if len(all_results) % 10 == 0:
+                        save_json_file(all_results, results_file)
+        
+        # Gracefully shut down the embedding server
+        print("All tasks completed. Shutting down the embedding server...")
+        request_queue.put((None, "STOP"))
+        embedding_server.join()
+
+    print("\n--- All Videos Processed ---")
+    save_json_file(all_results, results_file)
+    print(f"Detailed results saved to: {results_file}")
+
+    final_metrics = calculate_metrics(all_results)
+    save_json_file(final_metrics, metrics_file)
+    print(f"Final evaluation metrics saved to: {metrics_file}")
+    print(json.dumps(final_metrics, indent=4))
+
+
+if __name__ == "__main__":
+    # Before running, please ensure you have set the API Key for the image generation service
+    # export ARK_API_KEY="YOUR_VOLCENGINE_ARK_API_KEY"
+    main()
+

main_mcot.py

@@ -0,0 +1,608 @@
+import os
+import json
+import base64
+import argparse
+import time
+import re
+import traceback
+from datetime import datetime
+from functools import partial
+import requests  # Import requests library to download images from URLs
+from openai import AzureOpenAI, OpenAI
+from volcenginesdkarkruntime import Ark
+import concurrent.futures
+from tqdm import tqdm
+
+# 1. New Agent System Prompt
+# Defines the agent's role and principles, guiding it to use the "imagination" tool when visual evidence is insufficient.
+IMAGINE_AGENT_SYSTEM_PROMPT = """
+You are an intelligent AI assistant specializing in answering video question-answering problems through reasoning and imagination.
+Your task is to answer a multiple-choice question based on an initial, limited set of video frames.
+
+You will receive a few uniformly sampled frames to get a basic understanding of the video.
+These frames may not contain all the visual evidence needed to directly answer the question.
+
+If the provided frame information is insufficient, you must use the `imagine_frame` tool to generate new, imagined frames to fill in the visual gaps and aid your reasoning.
+You can call this tool multiple times to construct a sequence of imagined events.
+
+Your strategy should be:
+1. Analyze the initial frames and the user's question.
+2. Form a hypothesis about the missing content.
+3. If you need more visual information, call the `imagine_frame` tool. Provide a text `prompt` describing the scene you want to imagine, and select a `reference_image_id` from existing frames. The `reference_image_id` MUST be one of the IDs explicitly provided to you in the conversation history (e.g., "Frame ID: X" or "New Frame ID: Y"). Do not invent or assume frame IDs.
+4. Analyze the newly generated frame in conjunction with the existing ones.
+5. Continue this process of reasoning and imagination until you are confident in your answer. Please ensure you have found or created the relevant visual cues before answering the question.
+6. Each tool call can only generate one frame.
+
+IMPORTANT: Your text `prompt` for image generation must be safe and general. Avoid descriptions that could be interpreted as sensitive, harmful, or explicit to prevent generation failures.
+
+After your reasoning, provide the final answer in a JSON code block. The JSON object must contain a key "answer" with a value of one of 'A', 'B', 'C', or 'D'.
+
+Your output must strictly follow this format:
+<Your step-by-step reasoning process here, including why you chose to imagine a certain frame>
+```json
+{"answer": "X"}
+```
+Do not include any other text after the JSON code block.
+"""
+
+# 2. New Tool Schema for imagine_frame
+# Defines the interface, parameters, and description for the `imagine_frame` tool.
+IMAGINE_FRAME_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "imagine_frame",
+        "description": "When visual evidence is insufficient, generates a new image based on a text prompt and a reference image to help answer the question. Use it to imagine what might have happened between the provided frames.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "reference_image_id": {
+                    "type": "integer",
+                    "description": "The ID of an existing frame to use as a style and content reference. It can be one of the original frames or a previously generated one.",
+                },
+                "prompt": {
+                    "type": "string",
+                    "description": "A detailed text description of the frame you want to imagine and generate.",
+                },
+            },
+            "required": ["reference_image_id", "prompt"],
+        },
+    },
+}
+
+
+# 3. Implementation of the `imagine_frame` tool
+def imagine_frame(
+    reference_image_id: int,
+    prompt: str,
+    all_frame_paths: dict,
+    output_dir: str,
+    generation_count: int,
+):
+    """
+    Tool implementation: Calls an image generation model to create a new frame.
+
+    Args:
+        reference_image_id (int): The ID of the reference frame.
+        prompt (str): The text prompt for image generation.
+        all_frame_paths (dict): A dictionary containing IDs and paths of all currently available frames (original + generated).
+        output_dir (str): The directory to save the generated image.
+        generation_count (int): The current generation count, used for naming the file.
+
+    Returns:
+        str or None: The path of the newly generated image on success, otherwise None.
+    """
+    print(f"\n[Tool Call] Imagining new frame with prompt: '{prompt}'")
+    ark_api_key = os.environ.get("ARK_API_KEY")
+    if not ark_api_key:
+        raise ValueError("Error: Environment variable ARK_API_KEY is not set.")
+
+    client = Ark(
+        base_url="https://ark.cn-beijing.volces.com/api/v3",
+        api_key=ark_api_key,
+    )
+
+    ref_image_path = all_frame_paths.get(reference_image_id)
+    if not ref_image_path or not os.path.exists(ref_image_path):
+        raise FileNotFoundError(f"Reference image ID not found: {reference_image_id}")
+
+    try:
+        # Encode the reference image to a Base64 Data URI
+        ref_image_b64 = encode_image(ref_image_path)
+        ref_image_data_uri = f"data:image/jpeg;base64,{ref_image_b64}"
+
+        imagesResponse = client.images.generate(
+            model="doubao-seedream-4-0-250828",
+            prompt=prompt,
+            image=ref_image_data_uri,
+            size="1024x1024",  # Can be adjusted as needed, e.g., "2K"
+            response_format="url",
+            watermark=False,
+        )
+
+        image_url = imagesResponse.data[0].url
+
+        # Download the image from the URL
+        response = requests.get(image_url)
+        response.raise_for_status()
+
+        # Save the image to the specified directory
+        new_frame_filename = (
+            f"generated_frame_{generation_count}_ref_{reference_image_id}.jpg"
+        )
+        new_frame_path = os.path.join(output_dir, new_frame_filename)
+
+        with open(new_frame_path, "wb") as f:
+            f.write(response.content)
+
+        print(f"[Tool Success] Generated frame saved to: {new_frame_path}")
+        return new_frame_path
+
+    except Exception as e:
+        print(f"An error occurred during image generation or download: {e}")
+        traceback.print_exc()
+        return None
+
+
+def parse_arguments():
+    """Parse command-line arguments"""
+    parser = argparse.ArgumentParser(
+        description="Video QA Evaluation Framework with Imagine-and-Reason Agent"
+    )
+    parser.add_argument(
+        "--target-model",
+        "-tm",
+        type=str,
+        required=True,
+        help="The model to be evaluated (e.g., gpt-4o)",
+    )
+    parser.add_argument(
+        "--frames-path",
+        "-fp",
+        type=str,
+        required=True,
+        help="Absolute path to the root directory containing video frames.",
+    )
+    parser.add_argument(
+        "--output-path",
+        "-op",
+        type=str,
+        default="./generated_outputs",
+        help="Path to store generated images and results.",
+    )
+    parser.add_argument(
+        "--data-file",
+        "-df",
+        type=str,
+        required=True,
+        help="Absolute path to the evaluation dataset JSON file.",
+    )
+    parser.add_argument(
+        "--initial-frames-num",
+        "-ifn",
+        type=int,
+        default=8,
+        help="Number of initial uniformly sampled frames.",
+    )
+    parser.add_argument(
+        "--max-retry-times",
+        "-mr",
+        type=int,
+        default=10,
+        help="Maximum number of retries for failed API calls.",
+    )
+    parser.add_argument(
+        "--pool-processes",
+        "-pp",
+        type=int,
+        default=10,
+        help="Number of parallel processes.",
+    )
+    parser.add_argument(
+        "--base_url",
+        type=str,
+        required=True,
+        help="API Endpoint URL for the target model service.",
+    )
+    parser.add_argument(
+        "--api_key",
+        type=str,
+        required=True,
+        help="API Key for the target model service.",
+    )
+    return parser.parse_args()
+
+
+def save_json_file(data, output_file):
+    """Save data to a JSON file"""
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4, ensure_ascii=False)
+
+
+def extract_json_from_response(response):
+    """Extract JSON answer from the model's text response"""
+    if not response:
+        return None
+    match = re.search(r"```json\s*(\{.*?\})\s*```", response, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(1))
+        except (json.JSONDecodeError, IndexError):
+            return None
+    return None
+
+
+def calculate_metrics(results):
+    """Calculate various metrics from the evaluation results"""
+    valid_results = [r for r in results if "error" not in r]
+    total_samples = len(valid_results)
+    if total_samples == 0:
+        return {
+            "total_samples": 0,
+            "answered_samples": 0,
+            "correct_answers": 0,
+            "accuracy": 0.0,
+        }
+    answered_samples = sum(
+        1 for x in valid_results if x.get("model_answer") is not None
+    )
+    correct_answers = sum(1 for x in valid_results if x.get("is_correct"))
+    accuracy = correct_answers / answered_samples if answered_samples > 0 else 0.0
+    return {
+        "total_samples": total_samples,
+        "answered_samples": answered_samples,
+        "correct_answers": correct_answers,
+        "accuracy": accuracy,
+    }
+
+
+def call_single_model(client, messages, model, item_id, max_retry_times, tools=None):
+    """A single model API call with retry logic"""
+    params = {"model": model, "messages": messages, "max_tokens": 4096}
+    if tools:
+        params["tools"] = tools
+        params["tool_choice"] = "auto"
+
+    retry_times = 0
+    while retry_times < max_retry_times:
+        try:
+            completion = client.chat.completions.create(**params)
+            return completion.choices[0].message
+        except Exception as e:
+            retry_times += 1
+            print(
+                f"API call error (Item {item_id}): {str(e)}. Retrying ({retry_times}/{max_retry_times})..."
+            )
+            if retry_times == max_retry_times:
+                raise e
+            time.sleep(5)
+
+
+def uniformly_sample_frames_and_encode(frames_dir, num_frames):
+    """Uniformly sample a specified number of frames from a directory and encode them"""
+    if not os.path.isdir(frames_dir):
+        return [], {}
+
+    frame_files = sorted(
+        [f for f in os.listdir(frames_dir) if f.endswith(".jpg")],
+        key=lambda x: int(re.search(r"frame_(\d+)\.jpg", x).group(1)),
+    )
+
+    total_frames = len(frame_files)
+    if total_frames == 0:
+        return [], {}
+
+    if total_frames > num_frames:
+        indices = [int(i * total_frames / num_frames) for i in range(num_frames)]
+        sampled_files = [frame_files[i] for i in indices]
+    else:
+        sampled_files = frame_files
+
+    frame_path_map = {}
+    encoded_frames = []
+    for f in sampled_files:
+        path = os.path.join(frames_dir, f)
+        frame_id = int(re.search(r"frame_(\d+)\.jpg", f).group(1))
+        b64_image = encode_image(path)
+        # Send frame ID and image content as a pair
+        encoded_frames.append({"type": "text", "text": f"This is Frame ID: {frame_id}"})
+        encoded_frames.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/jpeg;base64,{b64_image}"},
+            }
+        )
+        frame_path_map[frame_id] = path
+
+    return encoded_frames, frame_path_map
+
+
+def evaluate_single_item_agentic_imagination(
+    data_item,
+    initial_frames,
+    initial_frame_paths,
+    generated_images_dir,
+    target_model,
+    api_key,
+    base_url,
+    max_retry_times,
+):
+    """
+    Core logic for evaluating a single data item using the Imagine-and-Reason Agent.
+    """
+    # 4. New Agent Loop
+    if "ark" in base_url:
+        client = Ark(base_url=base_url, api_key=api_key)
+    elif "aliyun" in base_url or "127.0.0.1" in base_url:
+        client = OpenAI(api_key=api_key, base_url=base_url)
+    else:
+        client = AzureOpenAI(
+            api_version="2023-05-15", api_key=api_key, azure_endpoint=base_url
+        )
+
+    tools = [IMAGINE_FRAME_TOOL_SCHEMA]
+
+    # Store paths of all available frames (initial + generated) in a dictionary for reference
+    available_frame_paths = initial_frame_paths.copy()
+
+    initial_prompt_content = [
+        {
+            "type": "text",
+            "text": "Here are the initial sampled video frames provided to you:",
+        },
+        *initial_frames,
+        {
+            "type": "text",
+            "text": f"Please answer the following question:\n{data_item['question']}",
+        },
+    ]
+
+    messages = [
+        {"role": "system", "content": IMAGINE_AGENT_SYSTEM_PROMPT},
+        {"role": "user", "content": initial_prompt_content},
+    ]
+
+    response_content = None
+    max_tool_calls = (
+        5  # Limit the number of times the agent can imagine to prevent infinite loops
+    )
+    generation_count = 0
+
+    for i in range(max_tool_calls):
+        response_message = call_single_model(
+            client,
+            messages,
+            target_model,
+            data_item["key"],
+            max_retry_times,
+            tools=tools,
+        )
+        if response_message is None:
+            return None
+
+        messages.append(response_message.model_dump(exclude_none=True))
+
+        if response_message.tool_calls:
+            tool_call = response_message.tool_calls[
+                0
+            ]  # Process one tool call at a time
+            function_name = tool_call.function.name
+
+            if function_name == "imagine_frame":
+                generation_count += 1
+                function_args = json.loads(tool_call.function.arguments)
+                new_frame_path = imagine_frame(
+                    **function_args,
+                    all_frame_paths=available_frame_paths,
+                    output_dir=generated_images_dir,
+                    generation_count=generation_count,
+                )
+
+                if new_frame_path:
+                    # Create a unique ID for the newly generated frame
+                    new_frame_id = (
+                        max(available_frame_paths.keys())
+                        if available_frame_paths
+                        else 0
+                    ) + 1
+                    available_frame_paths[new_frame_id] = new_frame_path
+
+                    b64_image = encode_image(new_frame_path)
+                    tool_response_content = [
+                        {
+                            "type": "text",
+                            "text": f"Here is the frame you requested to imagine (New Frame ID: {new_frame_id}). Please use it to continue your reasoning.",
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/jpeg;base64,{b64_image}"},
+                        },
+                    ]
+
+                    messages.append(
+                        {
+                            "tool_call_id": tool_call.id,
+                            "role": "tool",
+                            "name": function_name,
+                            "content": json.dumps(
+                                {"status": "success", "new_frame_id": new_frame_id}
+                            ),
+                        }
+                    )
+                    messages.append({"role": "user", "content": tool_response_content})
+                else:  # Tool execution failed
+                    messages.append(
+                        {
+                            "tool_call_id": tool_call.id,
+                            "role": "tool",
+                            "name": function_name,
+                            "content": json.dumps(
+                                {
+                                    "status": "error",
+                                    "message": "Failed to generate image.",
+                                }
+                            ),
+                        }
+                    )
+        else:  # No tool call means the model is ready to give a final answer
+            response_content = response_message.content
+            break
+
+    # If the max number of calls is reached without an answer, force a final response
+    if response_content is None and response_message:
+        final_prompt = "You have reached the maximum number of tool calls. Please provide a final answer in the specified JSON format based on the information you have gathered so far."
+        messages.append({"role": "user", "content": final_prompt})
+        final_response_message = call_single_model(
+            client, messages, target_model, data_item["key"], max_retry_times
+        )
+        if final_response_message:
+            messages.append(final_response_message.model_dump(exclude_none=True))
+            response_content = final_response_message.content
+
+    is_correct = False
+    model_answer_cleaned = None
+    parsed_json = extract_json_from_response(response_content)
+    if parsed_json and "answer" in parsed_json:
+        model_answer_cleaned = str(parsed_json["answer"]).strip().upper()
+        gold_answer = data_item["answer"].strip().upper()
+        if model_answer_cleaned == gold_answer:
+            is_correct = True
+
+    return {
+        **data_item,
+        "agent_conversation": messages,
+        "model_reasoning_and_answer": response_content,
+        "model_answer": model_answer_cleaned,
+        "is_correct": is_correct,
+        "generated_images_path": generated_images_dir,  # 5. Store the path to intermediate generated images
+    }
+
+
+def encode_image(image_path):
+    """Encode an image file to a Base64 string"""
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+def process_single_data(data_item, args):
+    """Worker function to process a single data item in parallel"""
+    item_key = data_item["key"]
+    try:
+        # Create a separate subfolder for each video's generated images
+        generated_images_dir = os.path.join(
+            args.output_path, "generated_images", item_key
+        )
+        os.makedirs(generated_images_dir, exist_ok=True)
+
+        specific_frames_path = os.path.join(args.frames_path, item_key)
+        initial_frames, initial_frame_paths = uniformly_sample_frames_and_encode(
+            specific_frames_path, args.initial_frames_num
+        )
+
+        if not initial_frames:
+            raise FileNotFoundError(f"Initial frames not found for item '{item_key}'")
+
+        result = evaluate_single_item_agentic_imagination(
+            data_item,
+            initial_frames,
+            initial_frame_paths,
+            generated_images_dir,
+            args.target_model,
+            args.api_key,
+            args.base_url,
+            args.max_retry_times,
+        )
+        return result
+
+    except Exception as e:
+        print(f"\nA critical error occurred while processing item {item_key}: {str(e)}")
+        traceback.print_exc()
+        return {
+            "key": item_key,
+            "uid": data_item.get("uid"),
+            "error": str(e),
+            "traceback": traceback.format_exc(),
+        }
+
+
+def load_test_data(json_file):
+    """Load test data from a JSON file"""
+    try:
+        with open(json_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        print(f"Error: Data file not found: {json_file}")
+        exit(1)
+    except json.JSONDecodeError:
+        print(f"Error: JSON file is malformed: {json_file}")
+        exit(1)
+
+
+def main():
+    """Main function to orchestrate the entire evaluation flow"""
+    args = parse_arguments()
+
+    print("--- Video QA Imagine-and-Reason Agent Framework ---")
+    print(f"Evaluating Model: {args.target_model}")
+    print(f"Output Path: {args.output_path}")
+    print(f"Dataset: {args.data_file}")
+    print("---------------------------------")
+
+    # Create the main output directory
+    os.makedirs(args.output_path, exist_ok=True)
+
+    model_name_safe = args.target_model.replace("/", "_")
+    data_filename_base = os.path.splitext(os.path.basename(args.data_file))[0]
+
+    output_prefix = f"{model_name_safe}_{data_filename_base}_imagine_agent"
+    results_output_file = os.path.join(
+        args.output_path, f"{output_prefix}_results.json"
+    )
+    metrics_output_file = os.path.join(
+        args.output_path, f"{output_prefix}_metrics.json"
+    )
+    error_log_file = os.path.join(args.output_path, f"{output_prefix}_errors.log")
+
+    # The logic for resuming from a checkpoint can be added here, same as in the first script
+
+    all_test_data = load_test_data(args.data_file)
+    tasks_to_process = all_test_data
+
+    all_results = []
+    # Use ProcessPoolExecutor for parallel processing
+    with concurrent.futures.ProcessPoolExecutor(
+        max_workers=args.pool_processes
+    ) as executor:
+        func = partial(process_single_data, args=args)
+        results_iterator = executor.map(func, tasks_to_process)
+
+        for result in tqdm(
+            results_iterator, total=len(tasks_to_process), desc="Processing Videos"
+        ):
+            if result:
+                if "error" in result:
+                    with open(error_log_file, "a", encoding="utf-8") as f:
+                        f.write(
+                            f"Error on item {result.get('key', 'N/A')}:\n  Error: {result['error']}\n---\n"
+                        )
+                all_results.append(result)
+
+                # Save results every 10 videos to prevent data loss from interruptions
+                if len(all_results) % 10 == 0:
+                    save_json_file(all_results, results_output_file)
+
+    print("\n\nProcessing complete.")
+    # Save the final complete results
+    save_json_file(all_results, results_output_file)
+    print(f"Detailed results saved to: {results_output_file}")
+
+    # Calculate and save the final metrics
+    final_metrics = calculate_metrics(all_results)
+    save_json_file(final_metrics, metrics_output_file)
+    print(f"\nEvaluation metrics saved to: {metrics_output_file}")
+    print(json.dumps(final_metrics, indent=4))
+
+
+if __name__ == "__main__":
+    # Before running this script, please ensure you have set the environment variable in your terminal:
+    # export ARK_API_KEY="YOUR_VOLCENGINE_ARK_API_KEY"
+    main()

main_new_agent.py

@@ -0,0 +1,643 @@
+import os
+import json
+import base64
+import argparse
+import time
+import re
+import traceback
+from datetime import datetime
+from functools import partial
+from openai import AzureOpenAI, OpenAI
+from volcenginesdkarkruntime import Ark
+import concurrent.futures
+from tqdm import tqdm
+import torch
+from transformers import AutoModel, AutoProcessor
+from torch.nn.functional import cosine_similarity
+# MODIFIED: Added imports for multiprocessing and uuid
+import multiprocessing
+import uuid
+
+# --- Configuration for SigLIP Model ---
+# MODIFIED: Updated to the local model path
+SIGLIP_MODEL_ID = (
+    "/mnt/bn/ziyang-storage-cloudnative-hl/huggingface/siglip-so400m-patch14-384"
+)
+
+# --- MODIFIED: Updated System Prompt explaining the two tools with examples ---
+AGENT_SYSTEM_PROMPT = """
+You are an intelligent AI assistant specialized in video question answering.
+Your task is to answer a multiple-choice question based on a video by strategically retrieving and analyzing its frames.
+
+You have two tools to retrieve frames. Both return images directly.
+
+1.  `get_frames_by_id(frame_ids)`: Retrieves frames using their specific numerical IDs. Use this when the question provides direct temporal clues or when you need to view specific frames identified by another tool.
+    * **Example Use Case:** For a question like "What happens at the 1 minute 30 second mark?", you can calculate the approximate frame ID and use this tool to see the visual.
+    * **Example Use Case:** For "Describe the action in frame 550.", you would call this tool with `frame_ids=[550]`.
+
+2.  `get_frames_by_similarity(query)`: Searches the entire video for frames that visually match a text description and returns the top 5 most relevant frames directly. Use this for content-based questions where the timing is unknown.
+    * **Example Use Case:** For a question like "What color is the main character's car?", you would use this tool with a query like "the main character's car".
+    * **Example Use Case:** For "Find the scene where a band is playing on stage", you would use the query "a band playing on stage".
+
+Your strategy must be efficient:
+1.  **Analyze the Query:** First, determine if the question is temporal/logical (better for `get_frames_by_id`) or content-based (requires `get_frames_by_similarity`).
+2.  **Retrieve & Analyze:** Call the most appropriate tool. Analyze the returned frames to form a hypothesis.
+3.  **Iterate:** If you need more information, refine your search query for the similarity tool or calculate new frame IDs for the ID tool and call again.
+4.  **Final Answer:** Once you have gathered enough visual evidence, provide your step-by-step reasoning and then the final answer in the specified JSON format. Do not guess.
+
+Your output should follow this format exactly:
+<Your step-by-step reasoning here>
+```json
+{"answer": "X"}
+```
+Do not include any other text after the JSON block.
+"""
+
+# Tool Schemas
+GET_FRAMES_BY_ID_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "get_frames_by_id",
+        "description": "Retrieves specific video frames by their numerical IDs to get visual information.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "frame_ids": {
+                    "type": "array",
+                    "items": {"type": "integer"},
+                    "description": "A list of up to 10 frame numbers to retrieve.",
+                },
+            },
+            "required": ["frame_ids"],
+        },
+    },
+}
+
+GET_FRAMES_BY_SIMILARITY_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "get_frames_by_similarity",
+        "description": "Searches for and retrieves the top 5 most visually relevant frames for a given text query. Use this to locate visual content when frame numbers are unknown.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "A concise text description of the visual content to search for (e.g., 'a person playing piano').",
+                },
+            },
+            "required": ["query"],
+        },
+    },
+}
+
+
+def parse_arguments():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Agentic Video QA with Hybrid Frame Retrieval"
+    )
+    parser.add_argument(
+        "--target-model", "-tm", type=str, required=True, help="Model to evaluate."
+    )
+    parser.add_argument(
+        "--frames-path",
+        "-fp",
+        type=str,
+        required=True,
+        help="Base directory for video frames.",
+    )
+    parser.add_argument(
+        "--data-file",
+        "-df",
+        type=str,
+        required=True,
+        help="Path to the evaluation dataset.",
+    )
+    parser.add_argument(
+        "--embeddings-path",
+        "-ep",
+        type=str,
+        required=True,
+        help="Directory with pre-computed frame embeddings.",
+    )
+    parser.add_argument(
+        "--max-retry-times",
+        "-mr",
+        type=int,
+        default=10,
+        help="Max retries for API calls.",
+    )
+    parser.add_argument(
+        "--pool-processes",
+        "-pp",
+        type=int,
+        default=20,
+        help="Number of parallel processes.",
+    )
+    parser.add_argument("--base_url", type=str, required=True, help="API endpoint URL.")
+    parser.add_argument("--api_key", type=str, required=True, help="API key.")
+    return parser.parse_args()
+
+
+def save_json_file(data, output_file):
+    """Saves data to a JSON file."""
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4)
+
+
+def extract_json_from_response(response):
+    """Extracts a JSON object from a model's response string."""
+    if not response:
+        return None
+    match = re.search(r"```json\s*(\{.*?\})\s*```", response, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(1))
+        except (json.JSONDecodeError, IndexError):
+            return None
+    return None
+
+
+def calculate_metrics(results):
+    """Calculates accuracy and other metrics from evaluation results."""
+    valid_results = [r for r in results if "error" not in r]
+    total_samples = len(valid_results)
+    if total_samples == 0:
+        return {
+            "total_samples": 0,
+            "answered_samples": 0,
+            "correct_answers": 0,
+            "accuracy": 0.0,
+        }
+    answered_samples = sum(
+        1 for x in valid_results if x.get("model_answer") is not None
+    )
+    correct_answers = sum(1 for x in valid_results if x.get("is_correct"))
+    accuracy = correct_answers / answered_samples if answered_samples > 0 else 0.0
+    return {
+        "total_samples": total_samples,
+        "answered_samples": answered_samples,
+        "correct_answers": correct_answers,
+        "accuracy": accuracy,
+    }
+
+
+def call_single_model(client, messages, model, item_id, max_retry_times, tools=None):
+    """Makes a single API call with retry logic and tool support."""
+    params = {"model": model, "messages": messages, "max_tokens": 4096}
+    if tools:
+        params["tools"] = tools
+        params["tool_choice"] = "auto"
+
+    for retry in range(max_retry_times):
+        try:
+            completion = client.chat.completions.create(**params)
+            return completion.choices[0].message
+        except Exception as e:
+            print(
+                f"API Error for item {item_id}: {str(e)}. Retrying ({retry + 1}/{max_retry_times})..."
+            )
+            if retry == max_retry_times - 1:
+                raise e
+            time.sleep(5)
+
+
+def get_frames_by_id(frame_ids: list, all_frame_paths: list):
+    """Tool implementation: Retrieves and encodes frames from a list of IDs."""
+    retrieved_frames = []
+    frame_map = {
+        int(re.search(r"frame_(\d+)\.jpg", os.path.basename(p)).group(1)): p
+        for p in all_frame_paths
+        if re.search(r"frame_(\d+)\.jpg", os.path.basename(p))
+    }
+    for fid in frame_ids:
+        path = frame_map.get(fid)
+        if path and os.path.exists(path):
+            b64_image = encode_image(path)
+            retrieved_frames.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/jpeg;base64,{b64_image}"},
+                }
+            )
+    return retrieved_frames
+
+
+# MODIFIED: This function is now the "client" side of the embedding service.
+def get_frames_by_similarity(
+    query: str,
+    all_frame_paths: list,
+    precomputed_data: dict,
+    request_queue: multiprocessing.Queue,
+    results_dict: dict,
+    k: int = 5,
+):
+    """
+    Requests a text embedding from the server process, calculates similarity,
+    finds top-k frames, and returns them encoded.
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    frame_filenames = precomputed_data["filenames"]
+    frame_embeddings = precomputed_data["embeddings"].to(device)
+
+    # 1. Send request to the embedding server process
+    request_id = str(uuid.uuid4())
+    request_queue.put((request_id, query))
+
+    # 2. Wait for the result
+    while request_id not in results_dict:
+        time.sleep(0.05)
+    query_embedding = results_dict.pop(request_id).to(device)
+
+    # 3. Perform similarity search with the received embedding
+    with torch.no_grad():
+        similarities = cosine_similarity(query_embedding, frame_embeddings)
+
+        num_frames_to_select = min(k, len(frame_filenames))
+        top_k_indices = (
+            torch.topk(similarities, k=num_frames_to_select, dim=-1)
+            .indices.cpu()
+            .flatten()
+            .numpy()
+        )
+
+        top_k_filenames = [frame_filenames[i] for i in top_k_indices]
+        top_k_frame_ids = [
+            int(re.search(r"frame_(\d+)\.jpg", f).group(1)) for f in top_k_filenames
+        ]
+
+    retrieved_frames = get_frames_by_id(top_k_frame_ids, all_frame_paths)
+    return retrieved_frames
+
+
+def evaluate_single_item_agentic(
+    data_item,
+    all_frame_paths,
+    embeddings_data,
+    target_model,
+    api_key,
+    base_url,
+    max_retry_times,
+    request_queue,  # MODIFIED: Added queue for IPC
+    results_dict,  # MODIFIED: Added dict for IPC
+):
+    """Evaluates a single item using an agentic loop."""
+    if "ark" in base_url:
+        client = Ark(base_url=base_url, api_key=api_key)
+    elif "aliyun" in base_url or "127.0.0.1" in base_url:
+        client = OpenAI(api_key=api_key, base_url=base_url)
+    else:
+        client = AzureOpenAI(
+            api_version="2023-05-15", api_key=api_key, azure_endpoint=base_url
+        )
+
+    tools = [GET_FRAMES_BY_ID_TOOL_SCHEMA, GET_FRAMES_BY_SIMILARITY_TOOL_SCHEMA]
+
+    get_frames_by_id_with_context = partial(
+        get_frames_by_id, all_frame_paths=all_frame_paths
+    )
+    # MODIFIED: Pass the request queue and results dict to the similarity function
+    get_frames_by_similarity_with_context = partial(
+        get_frames_by_similarity,
+        all_frame_paths=all_frame_paths,
+        precomputed_data=embeddings_data,
+        request_queue=request_queue,
+        results_dict=results_dict,
+    )
+
+    available_functions = {
+        "get_frames_by_id": get_frames_by_id_with_context,
+        "get_frames_by_similarity": get_frames_by_similarity_with_context,
+    }
+
+    total_frames = len(all_frame_paths)
+    duration = data_item.get("video_info", {}).get("duration_minutes", 0) * 60
+    initial_prompt = (
+        f"The video has {total_frames} frames (ID 1 to {total_frames}) and is {duration:.0f} seconds long. "
+        f"Please answer this question:\n{data_item['question']}"
+    )
+
+    messages = [
+        {"role": "system", "content": AGENT_SYSTEM_PROMPT},
+        {"role": "user", "content": initial_prompt},
+    ]
+    response_content = None
+    max_tool_calls = 10
+
+    for _ in range(max_tool_calls):
+        response_message = call_single_model(
+            client,
+            messages,
+            target_model,
+            data_item["key"],
+            max_retry_times,
+            tools=tools,
+        )
+        if response_message is None:
+            return None
+
+        messages.append(response_message)
+
+        if response_message.tool_calls:
+            for tool_call in response_message.tool_calls:
+                function_name = tool_call.function.name
+                function_to_call = available_functions.get(function_name)
+                if function_to_call:
+                    function_args = json.loads(tool_call.function.arguments)
+                    function_response = function_to_call(**function_args)
+
+                    messages.append(
+                        {
+                            "tool_call_id": tool_call.id,
+                            "role": "tool",
+                            "name": function_name,
+                            "content": json.dumps(
+                                {
+                                    "status": "success",
+                                    "retrieved_frame_count": len(function_response),
+                                }
+                            ),
+                        }
+                    )
+
+                    user_message_with_frames = [
+                        {
+                            "type": "text",
+                            "text": f"Here are the {len(function_response)} frames from your call to `{function_name}`.",
+                        }
+                    ]
+                    user_message_with_frames.extend(function_response)
+                    messages.append(
+                        {"role": "user", "content": user_message_with_frames}
+                    )
+        else:
+            response_content = response_message.content
+            break
+
+    if response_content is None:
+        final_prompt = "You have reached the maximum number of tool calls. Provide a final answer based on the information gathered so far."
+        messages.append({"role": "user", "content": final_prompt})
+        final_response = call_single_model(
+            client, messages, target_model, data_item["key"], max_retry_times
+        )
+        response_content = (
+            final_response.content
+            if final_response
+            else "Could not determine an answer after max tool calls."
+        )
+
+    is_correct = False
+    model_answer_cleaned = None
+    parsed_json = extract_json_from_response(response_content)
+    if parsed_json and "answer" in parsed_json:
+        model_answer_cleaned = str(parsed_json["answer"]).strip().upper()
+        if model_answer_cleaned == data_item["answer"].strip().upper():
+            is_correct = True
+
+    return {
+        **data_item,
+        "agent_conversation": [
+            msg if isinstance(msg, dict) else msg.model_dump() for msg in messages
+        ],
+        "model_reasoning_and_answer": response_content,
+        "model_answer": model_answer_cleaned,
+        "is_correct": is_correct,
+    }
+
+
+def encode_image(image_path):
+    """Encodes an image file to a base64 string."""
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+# MODIFIED: Function signature updated to accept queues for IPC
+def process_single_data(data_item, args, request_queue, results_dict):
+    """Main processing function for a single video, executed by a worker."""
+    item_key = data_item["key"]
+    try:
+        specific_frames_path = os.path.join(args.frames_path, item_key)
+        embedding_file = os.path.join(args.embeddings_path, f"{item_key}.pt")
+
+        if not os.path.isdir(specific_frames_path):
+            raise FileNotFoundError(
+                f"Frame directory not found: {specific_frames_path}"
+            )
+        if not os.path.exists(embedding_file):
+            raise FileNotFoundError(f"Embedding file not found: {embedding_file}")
+
+        all_frame_paths = sorted(
+            [
+                os.path.join(specific_frames_path, f)
+                for f in os.listdir(specific_frames_path)
+                if f.endswith(".jpg")
+            ],
+            key=lambda x: int(re.search(r"frame_(\d+)\.jpg", x).group(1)),
+        )
+        if not all_frame_paths:
+            raise FileNotFoundError(f"No frames found for key '{item_key}'")
+
+        embeddings_data = torch.load(embedding_file, map_location="cpu")
+
+        # MODIFIED: Pass queues to the evaluation function
+        result = evaluate_single_item_agentic(
+            data_item,
+            all_frame_paths,
+            embeddings_data,
+            args.target_model,
+            args.api_key,
+            args.base_url,
+            args.max_retry_times,
+            request_queue,
+            results_dict,
+        )
+        return result
+
+    except Exception as e:
+        print(f"\nCRITICAL ERROR on key {item_key}: {str(e)}")
+        traceback.print_exc()
+        return {
+            "key": item_key,
+            "uid": data_item.get("uid"),
+            "error": str(e),
+            "traceback": traceback.format_exc(),
+        }
+
+
+def load_test_data(json_file):
+    """Loads the evaluation data from a JSON file."""
+    try:
+        with open(json_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        print(f"Error: Data file not found: {json_file}")
+        exit(1)
+    except json.JSONDecodeError:
+        print(f"Error: Malformed JSON in {json_file}")
+        exit(1)
+
+
+# MODIFIED: This new function runs in its own process, handling embedding requests.
+# It now accepts a model_id and loads the model itself.
+def embedding_server_process(model_id, device, request_queue, results_dict):
+    """
+    A server process that loads the SigLIP model and continuously fetches
+    text queries from a queue, computes their embeddings, and places the
+    results in a shared dictionary.
+    """
+    print(f"Embedding server started on PID {os.getpid()}...")
+    print("Loading SigLIP model in the embedding server process...")
+    model = AutoModel.from_pretrained(model_id)
+    processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
+    print("SigLIP model loaded in server.")
+
+    model.to(device)
+    model.eval()
+
+    while True:
+        try:
+            request_id, text_query = request_queue.get()
+            if text_query == "STOP":
+                print("Embedding server received stop signal. Shutting down.")
+                break
+
+            with torch.no_grad():
+                text_inputs = processor(
+                    text=[text_query],
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                ).to(device)
+                query_embedding = model.get_text_features(**text_inputs)
+                # Move embedding to CPU before sharing across processes
+                results_dict[request_id] = query_embedding.cpu()
+        except Exception as e:
+            print(f"Error in embedding server: {e}")
+            traceback.print_exc()
+
+
+# MODIFIED: The old init_worker is removed.
+def main():
+    """Main function to orchestrate the evaluation framework."""
+    args = parse_arguments()
+    print("--- Agentic Video QA with Hybrid Retrieval ---")
+    print(
+        f"Model: {args.target_model}, Data: {args.data_file}, Embeddings: {args.embeddings_path}"
+    )
+
+    # MODIFIED: Changed start method to 'spawn' for safety with CUDA and on macOS/Windows.
+    try:
+        multiprocessing.set_start_method("spawn", force=True)
+        print("Multiprocessing start method set to 'spawn'.")
+    except RuntimeError:
+        print("Start method already set.")
+
+    # MODIFIED: Model is no longer loaded in the main process.
+    # It will be loaded in the dedicated embedding_server_process.
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model_name_safe = args.target_model.replace("/", "_")
+    data_filename_base = os.path.splitext(os.path.basename(args.data_file))[0]
+    output_prefix = f"{model_name_safe}_{data_filename_base}_agent_hybrid"
+    results_output_file = f"{output_prefix}_results.json"
+    metrics_output_file = f"{output_prefix}_metrics.json"
+    error_log_file = f"{output_prefix}_errors.log"
+
+    with open(error_log_file, "a", encoding="utf-8") as f:
+        f.write(
+            f"\n=== Log Session Started at {datetime.now()} for {args.target_model} ===\n"
+        )
+
+    all_test_data = load_test_data(args.data_file)
+    existing_results = []
+    completed_ids = set()
+    if os.path.exists(results_output_file):
+        try:
+            with open(results_output_file, "r", encoding="utf-8") as f:
+                existing_results = json.load(f)
+            if isinstance(existing_results, list):
+                completed_ids = {
+                    item["uid"] for item in existing_results if "uid" in item
+                }
+                print(f"Found {len(completed_ids)} completed tasks. Resuming...")
+            else:
+                existing_results = []
+        except (json.JSONDecodeError, IOError):
+            existing_results = []
+
+    tasks_to_process = [
+        item for item in all_test_data if item.get("uid") not in completed_ids
+    ]
+    if not tasks_to_process:
+        print("All tasks are already completed. Calculating final metrics.")
+    else:
+        print(
+            f"Total: {len(all_test_data)}. Completed: {len(completed_ids)}. To process: {len(tasks_to_process)}."
+        )
+
+    all_results = list(existing_results)
+
+    if tasks_to_process:
+        # MODIFIED: Set up Manager, Queues, and the embedding server process
+        with multiprocessing.Manager() as manager:
+            request_queue = manager.Queue()
+            results_dict = manager.dict()
+
+            # MODIFIED: Start the dedicated embedding server process, passing the model ID.
+            embedding_server = multiprocessing.Process(
+                target=embedding_server_process,
+                args=(
+                    SIGLIP_MODEL_ID,
+                    device,
+                    request_queue,
+                    results_dict,
+                ),
+            )
+            embedding_server.start()
+
+            # MODIFIED: The ProcessPoolExecutor no longer needs an initializer for the model
+            with concurrent.futures.ProcessPoolExecutor(
+                max_workers=args.pool_processes
+            ) as executor:
+                # MODIFIED: Pass the queues to each worker via partial
+                func = partial(
+                    process_single_data,
+                    args=args,
+                    request_queue=request_queue,
+                    results_dict=results_dict,
+                )
+                results_iterator = executor.map(func, tasks_to_process)
+                for result in tqdm(
+                    results_iterator,
+                    total=len(tasks_to_process),
+                    desc="Processing Videos",
+                ):
+                    if result:
+                        if "error" in result:
+                            with open(error_log_file, "a", encoding="utf-8") as f:
+                                f.write(
+                                    f"Error on key {result.get('key', 'N/A')}:\n  Error: {result['error']}\n  Traceback: {result['traceback']}\n---\n"
+                                )
+                        all_results.append(result)
+                        if len(all_results) % 10 == 0:
+                            save_json_file(all_results, results_output_file)
+
+            # MODIFIED: Gracefully shut down the embedding server
+            print("All tasks processed. Sending stop signal to embedding server.")
+            request_queue.put((None, "STOP"))
+            embedding_server.join()
+
+    print("\n\nProcessing complete.")
+    save_json_file(all_results, results_output_file)
+    print(f"Detailed results saved to: {results_output_file}")
+
+    final_metrics = calculate_metrics(all_results)
+    save_json_file(final_metrics, metrics_output_file)
+    print(f"\nMetrics saved to: {metrics_output_file}")
+    print(json.dumps(final_metrics, indent=4))
+
+
+if __name__ == "__main__":
+    main()
+

main_uniform_sampling.py

@@ -0,0 +1,501 @@
+import os
+import json
+import base64
+import argparse
+import time
+import re
+from datetime import datetime
+from functools import partial
+from openai import AzureOpenAI, OpenAI
+from volcenginesdkarkruntime import Ark
+from multiprocessing import Pool, Manager, Lock
+
+# New prompt template for multiple-choice questions with reasoning
+REASONING_MULTIPLE_CHOICE_TEMPLATE = """
+You are an AI assistant evaluating video frames to answer a multiple-choice question.
+The user will provide you with a set of video frames and a question with several options (e.g., A, B, C, D).
+
+First, provide a step-by-step reasoning process that analyzes the video frames and leads to your conclusion.
+After your reasoning, provide the final answer in a JSON block. The JSON object must contain a single key "answer" with the value being one of 'A', 'B', 'C', or 'D'.
+
+Your output should follow this format exactly:
+<Your step-by-step reasoning here>
+```json
+{"answer": "A"}
+```
+Do not include any other text after the JSON block.
+"""
+
+
+def parse_arguments():
+    """
+    Parse command line arguments for evaluation configuration.
+
+    Returns:
+        argparse.Namespace: Parsed command line arguments
+    """
+    parser = argparse.ArgumentParser(description="Video QA Evaluation Framework")
+
+    # Model configuration
+    parser.add_argument(
+        "--target-model",
+        "-tm",
+        type=str,
+        required=True,
+        help="Model to be evaluated (e.g., gpt-4o, gpt-4-vision-preview)",
+    )
+
+    # Data configuration
+    parser.add_argument(
+        "--frame-num",
+        "-fn",
+        type=int,
+        default=32,
+        help="Number of frames to uniformly sample from each video (default: 32)",
+    )
+    parser.add_argument(
+        "--frames-path",
+        "-fp",
+        type=str,
+        required=True,
+        help="Absolute path to the base directory containing video frame folders.",
+    )
+    parser.add_argument(
+        "--data-file",
+        "-df",
+        type=str,
+        required=True,
+        help="Absolute path to the JSON file containing the evaluation dataset.",
+    )
+
+    # Processing configuration
+    parser.add_argument(
+        "--max-retry-times",
+        "-mr",
+        type=int,
+        default=10,
+        help="Maximum number of retries for API calls (default: 10)",
+    )
+    parser.add_argument(
+        "--pool-processes",
+        "-pp",
+        type=int,
+        default=20,
+        help="Number of parallel processes for evaluation (default: 20)",
+    )
+
+    # API configuration
+    parser.add_argument(
+        "--base_url", type=str, required=True, help="Azure OpenAI endpoint URL."
+    )
+    parser.add_argument(
+        "--api_key", type=str, required=True, help="Azure OpenAI API key."
+    )
+
+    return parser.parse_args()
+
+
+def save_json_file(data, output_file):
+    """
+    Save data to a JSON file.
+
+    Args:
+        data (dict or list): Data to be saved.
+        output_file (str): Path to the output file.
+    """
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4)
+
+
+def extract_json_from_response(response):
+    """
+    Extracts a JSON object from a string that contains reasoning followed by a tagged JSON block.
+
+    Args:
+        response (str): The raw response string from the model.
+
+    Returns:
+        dict or None: Parsed JSON object or None if no valid JSON block is found.
+    """
+    if not response:
+        return None
+    try:
+        # Regex to find the content inside ```json ... ```
+        match = re.search(r"```json\s*(\{.*?\})\s*```", response, re.DOTALL)
+        if match:
+            json_str = match.group(1)
+            return json.loads(json_str)
+        return None
+    except (json.JSONDecodeError, IndexError):
+        return None
+
+
+def calculate_metrics(results):
+    """
+    Calculate evaluation metrics from the results.
+
+    Args:
+        results (list): List of results with 'is_correct' field.
+
+    Returns:
+        dict: Dictionary containing calculated metrics.
+    """
+    total_samples = len(results)
+    if total_samples == 0:
+        return {
+            "total_samples": 0,
+            "answered_samples": 0,
+            "correct_answers": 0,
+            "accuracy": 0.0,
+        }
+
+    answered_samples = sum(1 for x in results if x.get("model_answer") is not None)
+    correct_answers = sum(1 for x in results if x.get("is_correct"))
+
+    accuracy = correct_answers / answered_samples if answered_samples > 0 else 0.0
+
+    metrics = {
+        "total_samples": total_samples,
+        "answered_samples": answered_samples,
+        "correct_answers": correct_answers,
+        "accuracy": accuracy,
+    }
+
+    return metrics
+
+
+def call_single_model(client, messages, model, item_id, max_retry_times):
+    """
+    Make a single API call to the specified model with retry logic.
+
+    Args:
+        client: OpenAI client instance.
+        messages (list): List of messages for the API call.
+        model (str): Model name to use.
+        item_id (str): ID of the item being processed (for error logging).
+        max_retry_times (int): Maximum number of retries.
+
+    Returns:
+        str or None: Model response or None if all retries failed.
+    """
+    if "doubao" in model:
+        max_tokens = 32768
+    else:
+        max_tokens = 65535
+    retry_times = 0
+    while retry_times < max_retry_times:
+        try:
+            # Set max_tokens to a larger value to allow for reasoning
+            completion = client.chat.completions.create(
+                model=model, messages=messages, max_tokens=max_tokens
+            )
+            return completion.choices[0].message.content
+        except Exception as e:
+            retry_times += 1
+            print(
+                f"Error processing item {item_id} with model {model}: {str(e)}. Retrying ({retry_times}/{max_retry_times})..."
+            )
+            if retry_times == max_retry_times:
+                error_log_file = f"error_log_{model.replace('/', '_')}.txt"
+                with open(error_log_file, "a") as f:
+                    f.write(
+                        f"Error processing item {item_id} with model {model} after {max_retry_times} retries: {str(e)}\n"
+                    )
+                return None
+            time.sleep(5)  # Wait before retrying
+
+
+def evaluate_single_item(
+    data_item, frames, target_model, api_key, base_url, max_retry_times
+):
+    """
+    Evaluate a single data item using the target model and perform exact match.
+
+    Args:
+        data_item (dict): Dictionary containing question and answer data.
+        frames (list): List of encoded video frames.
+        target_model (str): Model to be evaluated.
+        api_key (str): API key.
+        base_url (str): API base URL.
+        max_retry_times (int): Maximum number of retries.
+
+    Returns:
+        dict: Evaluation result.
+    """
+    if "ark" in base_url:
+        client = Ark(
+            base_url=base_url,
+            api_key=api_key,
+        )
+    elif "aliyun" in base_url or "127.0.0.1" in base_url:
+        client = OpenAI(api_key=api_key, base_url=base_url)
+    else:
+        client = AzureOpenAI(
+            api_version="2023-05-15", api_key=api_key, azure_endpoint=base_url
+        )
+
+    # Construct messages for the model using the new template
+    messages = [
+        {"role": "system", "content": REASONING_MULTIPLE_CHOICE_TEMPLATE},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Here are the video frames:"},
+                *frames,
+                {"type": "text", "text": f"Question: {data_item['question']}"},
+            ],
+        },
+    ]
+
+    response = call_single_model(
+        client, messages, target_model, data_item["key"], max_retry_times
+    )
+
+    is_correct = False
+    model_answer_cleaned = None
+    parsed_json = None
+
+    if response:
+        parsed_json = extract_json_from_response(response)
+        if parsed_json and "answer" in parsed_json:
+            model_answer_cleaned = str(parsed_json["answer"]).strip().upper()
+            gold_answer = data_item["answer"].strip().upper()
+            if model_answer_cleaned == gold_answer:
+                is_correct = True
+
+    # Create result dictionary
+    result = {
+        **data_item,
+        "model_reasoning_and_answer": response,
+        "model_answer_raw": parsed_json.get("answer") if parsed_json else None,
+        "model_answer": model_answer_cleaned,
+        "is_correct": is_correct,
+    }
+
+    return result
+
+
+def encode_image(image_path):
+    """
+    Encode an image file to base64 string.
+
+    Args:
+        image_path (str): Path to the image file.
+
+    Returns:
+        str: Base64 encoded image string.
+    """
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+
+def process_frames(frames_path, frame_num):
+    """
+    Process and uniformly sample video frames from a directory, then encode them.
+
+    Args:
+        frames_path (str): Path to the directory containing video frames.
+        frame_num (int): The number of frames to sample.
+
+    Returns:
+        list: List of encoded frame objects for API consumption.
+    """
+    if not os.path.isdir(frames_path):
+        print(f"Warning: Frame directory not found at {frames_path}")
+        return []
+
+    frame_files = [
+        f
+        for f in os.listdir(frames_path)
+        if f.startswith("frame_") and f.endswith(".jpg")
+    ]
+    # Sort frames numerically based on the ID in frame_{id}.jpg
+    frame_files.sort(key=lambda x: int(x.split("_")[1].split(".")[0]))
+
+    frame_path_list = [os.path.join(frames_path, f) for f in frame_files]
+    total_frames = len(frame_path_list)
+
+    if total_frames == 0:
+        return []
+
+    # Uniformly sample frame paths
+    if total_frames > frame_num:
+        indices = [int(i * total_frames / frame_num) for i in range(frame_num)]
+        sampled_paths = [frame_path_list[i] for i in indices]
+    else:
+        sampled_paths = frame_path_list  # Use all frames if fewer than requested
+
+    # Encode only the sampled frames
+    base64_images = [encode_image(path) for path in sampled_paths]
+
+    # Create frame objects for API payload
+    return [
+        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_img}"}}
+        for b64_img in base64_images
+    ]
+
+
+def process_single_data(
+    data_item, args, shared_results, progress_counter, total_items, locks
+):
+    """
+    Process a single data item in a multiprocessing context.
+
+    Args:
+        data_item (dict): Single data item to process.
+        args: Command line arguments.
+        shared_results: Shared list for storing results.
+        progress_counter: Shared counter for progress tracking.
+        total_items (int): Total number of items to process.
+        locks (dict): Dictionary of locks for thread-safe operations.
+    """
+    item_key = data_item["key"]
+    try:
+        # Construct path to the specific video's frames folder
+        specific_frames_path = os.path.join(args.frames_path, item_key)
+        frames = process_frames(specific_frames_path, args.frame_num)
+
+        if not frames:
+            raise FileNotFoundError(
+                f"No frames found or processed for key '{item_key}' at path '{specific_frames_path}'"
+            )
+
+        result = evaluate_single_item(
+            data_item,
+            frames,
+            args.target_model,
+            args.api_key,
+            args.base_url,
+            args.max_retry_times,
+        )
+
+        if result is not None:
+            with locks["results"]:
+                shared_results.append(result)
+                # Define output file names inside the worker
+                data_filename_base = os.path.splitext(os.path.basename(args.data_file))[
+                    0
+                ]
+                model_name_safe = args.target_model.replace("/", "_")
+                output_prefix = (
+                    f"{model_name_safe}_{data_filename_base}_{args.frame_num}frames"
+                )
+                results_output_file = f"{output_prefix}_results.json"
+                # Save the entire updated list of results after each case is processed
+                save_json_file(list(shared_results), results_output_file)
+
+    except Exception as e:
+        print(f"Error processing video key {item_key}: {str(e)}")
+        with locks["file"]:
+            error_log_file = f"error_log_{args.target_model.replace('/', '_')}.txt"
+            with open(error_log_file, "a") as f:
+                f.write(f"Critical error processing video key {item_key}: {str(e)}\n")
+    finally:
+        # Always update progress counter
+        with locks["counter"]:
+            progress_counter.value += 1
+            print(
+                f"\rProcessed: {progress_counter.value}/{total_items} videos...",
+                end="",
+                flush=True,
+            )
+
+
+def load_test_data(json_file):
+    """
+    Load test data from a JSON file.
+
+    Args:
+        json_file (str): Path to the JSON file.
+
+    Returns:
+        list: List of test data items.
+    """
+    try:
+        with open(json_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        print(f"Error: Data file not found at {json_file}")
+        exit(1)
+    except json.JSONDecodeError:
+        print(f"Error: Could not decode JSON from {json_file}")
+        exit(1)
+
+
+def main():
+    """
+    Main function to run the video QA evaluation framework.
+    """
+    args = parse_arguments()
+
+    print("--- Evaluation Configuration ---")
+    print(f"Target Model: {args.target_model}")
+    print(f"Frames to Sample: {args.frame_num}")
+    print(f"Frames Base Path: {args.frames_path}")
+    print(f"Data File: {args.data_file}")
+    print(f"Parallel Processes: {args.pool_processes}")
+    print("---------------------------------")
+
+    # Initialize error log file
+    error_log_file = f"error_log_{args.target_model.replace('/', '_')}.txt"
+    with open(error_log_file, "w") as f:
+        f.write(
+            f"=== Error Log Started at {datetime.now()} for model {args.target_model} ===\n"
+        )
+
+    # Define output file names
+    data_filename_base = os.path.splitext(os.path.basename(args.data_file))[0]
+    model_name_safe = args.target_model.replace("/", "_")
+    output_prefix = f"{model_name_safe}_{data_filename_base}_{args.frame_num}frames"
+
+    results_output_file = f"{output_prefix}_results.json"
+    metrics_output_file = f"{output_prefix}_metrics.json"
+
+    # Load data
+    test_data = load_test_data(args.data_file)
+    total_videos = len(test_data)
+    print(f"\nLoaded {total_videos} videos to process.")
+
+    # Set up multiprocessing
+    with Manager() as manager:
+        shared_results = manager.list()
+        progress_counter = manager.Value("i", 0)
+
+        locks = {
+            "results": manager.Lock(),
+            "file": manager.Lock(),
+            "counter": manager.Lock(),
+        }
+
+        # Create a partial function with fixed arguments for the worker pool
+        process_func = partial(
+            process_single_data,
+            args=args,
+            shared_results=shared_results,
+            progress_counter=progress_counter,
+            total_items=total_videos,
+            locks=locks,
+        )
+
+        # Run processing in parallel
+        with Pool(processes=args.pool_processes) as pool:
+            pool.map(process_func, test_data)
+
+        # Convert shared list to a regular list for final processing
+        all_results = list(shared_results)
+
+    print(f"\n\nProcessing complete for model: {args.target_model}")
+
+    # Calculate and save final metrics
+    final_metrics = calculate_metrics(all_results)
+    save_json_file(final_metrics, metrics_output_file)
+    print(f"\nMetrics saved to: {metrics_output_file}")
+    print(json.dumps(final_metrics, indent=4))
+
+    # Save final results
+    save_json_file(all_results, results_output_file)
+    print(f"Detailed results saved to: {results_output_file}")
+
+
+if __name__ == "__main__":
+    main()

offline_compute_similarity.py

@@ -0,0 +1,191 @@
+import os
+import json
+import torch
+import argparse
+import tempfile
+import glob
+from tqdm import tqdm
+from transformers import AutoModel, AutoProcessor
+from torch.nn.functional import cosine_similarity
+import torch.multiprocessing as mp
+
+# --- 配置 ---
+MODEL_ID = "/mnt/bn/ziyang-storage-cloudnative-hl/huggingface/siglip-so400m-patch14-384"
+
+
+def parse_arguments():
+    """解析命令行参数"""
+    parser = argparse.ArgumentParser(
+        description="步骤 2: 从预计算的嵌入加载并计算问-帧相似度。"
+    )
+    parser.add_argument(
+        "--data-file",
+        "-df",
+        type=str,
+        required=True,
+        help="包含评估数据集的JSON文件的绝对路径。",
+    )
+    parser.add_argument(
+        "--embeddings-path",
+        "-ep",
+        type=str,
+        required=True,
+        help="包含预计算嵌入.pt文件的目录的绝对路径。",
+    )
+    parser.add_argument(
+        "--output-file",
+        "-o",
+        type=str,
+        required=True,
+        help="用于保存最终相似度分数的JSON文件路径。",
+    )
+    return parser.parse_args()
+
+
+def load_test_data(json_file):
+    """从JSON文件加载测试数据"""
+    try:
+        with open(json_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        print(f"错误: 在 {json_file} 未找到数据文件")
+        exit(1)
+    except json.JSONDecodeError:
+        print(f"错误: 无法从 {json_file} 解码JSON")
+        exit(1)
+    return []
+
+
+def save_json_file(data, output_file):
+    """将数据保存到JSON文件"""
+    # os.makedirs(os.path.dirname(output_file), exist_ok=True)
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4)
+    print(f"\n成功将最终相似度结果保存到 {output_file}")
+
+
+def process_question_chunk(args_tuple):
+    """
+    工作函数，用于处理一批问题并增量保存结果。
+    """
+    data_chunk, embeddings_base_path, gpu_id, temp_dir = args_tuple
+    device = f"cuda:{gpu_id}"
+
+    # 为此工作进程定义一个唯一的临时输出文件
+    temp_output_file = os.path.join(temp_dir, f"results_gpu_{gpu_id}.jsonl")
+
+    # 只需要模型来计算文本特征
+    model = AutoModel.from_pretrained(MODEL_ID).to(device).eval()
+    processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=True)
+
+    progress_bar = tqdm(data_chunk, position=gpu_id, desc=f"GPU-{gpu_id}")
+
+    # 缓存已加载的嵌入以避免重复IO
+    embedding_cache = {}
+
+    with open(temp_output_file, "a", encoding="utf-8") as f_out:
+        for data_item in progress_bar:
+            item_key = data_item["key"]
+            question_key = data_item["uid"]
+            question = data_item["question"].split("\n(A)")[0]
+
+            embedding_file_path = os.path.join(embeddings_base_path, f"{item_key}.pt")
+            if not os.path.exists(embedding_file_path):
+                progress_bar.write(
+                    f"Warning: Embedding file not found for '{item_key}', skipping."
+                )
+                continue
+
+            try:
+                # 从缓存或文件中加载嵌入
+                if item_key not in embedding_cache:
+                    loaded_data = torch.load(embedding_file_path, map_location="cpu")
+                    embedding_cache[item_key] = {
+                        "filenames": loaded_data["filenames"],
+                        "embeddings": loaded_data["embeddings"],
+                    }
+
+                frame_files = embedding_cache[item_key]["filenames"]
+                frame_embeddings = embedding_cache[item_key]["embeddings"].to(device)
+
+                with torch.no_grad():
+                    # --- 文本嵌入 ---
+                    text_inputs = processor(
+                        text=[question],
+                        return_tensors="pt",
+                        padding=True,
+                        truncation=True,
+                    ).to(device)
+                    question_embedding = model.get_text_features(**text_inputs)
+
+                    # --- 相似度计算和排序 ---
+                    similarities = cosine_similarity(
+                        question_embedding, frame_embeddings
+                    )
+                    scored_frames = sorted(
+                        zip(frame_files, similarities.cpu().numpy()),
+                        key=lambda x: x[1],
+                        reverse=True,
+                    )
+                    sorted_frame_filenames = [frame[0] for frame in scored_frames]
+
+                    single_result = {question_key: sorted_frame_filenames}
+                    f_out.write(json.dumps(single_result) + "\n")
+
+            except Exception as e:
+                progress_bar.write(f"Error on GPU-{gpu_id} for item '{item_key}': {e}")
+
+
+def main():
+    """主函数，用于协调多GPU处理"""
+    args = parse_arguments()
+
+    num_gpus = torch.cuda.device_count()
+    if num_gpus == 0:
+        print("错误: 未找到启用CUDA的GPU。正在退出。")
+        exit(1)
+
+    print(f"找到 {num_gpus} 个GPU。开始并行计算相似度...")
+
+    test_data = load_test_data(args.data_file)
+    if not test_data:
+        return
+
+    chunk_size = (len(test_data) + num_gpus - 1) // num_gpus
+    data_chunks = [
+        test_data[i : i + chunk_size] for i in range(0, len(test_data), chunk_size)
+    ]
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        print(f"使用临时目录存储中间结果: {temp_dir}")
+
+        process_args = [
+            (data_chunks[i], args.embeddings_path, i, temp_dir)
+            for i in range(len(data_chunks))
+        ]
+
+        with mp.Pool(processes=num_gpus) as pool:
+            pool.map(process_question_chunk, process_args)
+
+        # --- 从临时文件合并和保存最终结果 ---
+        print("\n\n所有GPU进程已完成。正在从临时文件合并结果...")
+        final_similarity_results = {}
+
+        temp_files = glob.glob(os.path.join(temp_dir, "*.jsonl"))
+
+        for temp_file in tqdm(temp_files, desc="合并文件"):
+            with open(temp_file, "r", encoding="utf-8") as f:
+                for line in f:
+                    try:
+                        data = json.loads(line)
+                        final_similarity_results.update(data)
+                    except json.JSONDecodeError:
+                        print(f"警告: 跳过 {temp_file} 中的损坏行")
+
+        save_json_file(final_similarity_results, args.output_file)
+        print(f"总共处理的项目数: {len(final_similarity_results)}")
+
+
+if __name__ == "__main__":
+    mp.set_start_method("spawn", force=True)
+    main()

utils/count_frames.py

@@ -0,0 +1,35 @@
+import json
+
+
+with open(
+    "/mnt/bn/ziyang-storage-cloudnative-hl/VideoSimpleQA/o4-mini-2025-04-16_lvbench_agent_results.json",
+    "r",
+) as f:
+    data = json.load(f)
+
+# print(data[0])
+# import pdb; pdb.set_trace()
+# print(data[1])
+
+frames_count = []
+# for d in data:
+#     num_frames = 0
+#     for turn in d["agent_conversation"]:
+#         if turn["role"] == "assistant" and turn["tool_calls"]:
+#             for tool_call in turn["tool_calls"]:
+#                 frame_ids = json.loads(tool_call["function"]["arguments"])["frame_ids"]
+#                 num_frames += len(frame_ids)
+#     frames_count.append(num_frames)
+
+for d in data:
+    num_frames = 0
+    for turn in d["agent_conversation"]:
+        if turn["role"] == "user" and type(turn['content']) == list:
+            for item in turn['content']:
+                if item['type'] == 'image_url':
+                    num_frames += 1
+    frames_count.append(num_frames)
+
+print(f"mean frames: {sum(frames_count) / len(frames_count)}")
+print(f"max frames: {max(frames_count)}")
+print(f"min frames: {min(frames_count)}")