inference/video_dr_bridge.py

"""
VideoDR 评测桥接模块。

输入:
- VideoDR CSV 标注文件、原始视频目录，以及 demo 本地打包的 `video_dr_gen`
  工具实现。
- 评测阶段生成的 bbox、检索 query、本地 MARS 服务地址与图片搜索配置。

处理:
- 复用 SFT 构造代码里的统一 system prompt、bbox 归一化、图片搜索与 MARS web search 实现。
- 为当前仓库补充 VideoDR CSV 解析、原始视频按 1fps 抽帧、区间采样与单帧定位能力。
- 将 VideoDR 的无表头 CSV 行正规化为评测样本，并把视频路径映射到 `video/<id>.mp4`。

输出:
- 导出 VideoDR 统一 prompt、CSV loader、1fps 抽帧函数，以及评测器需要调用的工具函数。
- 为 `inference/eval.py` 提供与 SFT 构造代码一致的工具后端。
"""

from __future__ import annotations

import csv
import os
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np


# VideoDR SFT 构造代码（assemble_sft_dataset / config / prompts / utils）的位置。
# 解析顺序：环境变量 VIDEO_DR_SOURCE_ROOT > 本目录同级的打包目录 demo/video_dr_gen。
def _resolve_video_dr_source_root() -> Path:
    env_root = os.environ.get("VIDEO_DR_SOURCE_ROOT", "").strip()
    candidates = []
    if env_root:
        candidates.append(Path(env_root))
    # demo/inference/video_dr_bridge.py -> demo/video_dr_gen
    candidates.append(Path(__file__).resolve().parent.parent / "video_dr_gen")
    for cand in candidates:
        if cand.exists():
            return cand
    raise RuntimeError(
        "找不到 VideoDR SFT 构造代码目录（assemble_sft_dataset/config/prompts/utils）。"
        f" 已尝试: {[str(c) for c in candidates]}。"
        " 请设置环境变量 VIDEO_DR_SOURCE_ROOT，或把 video_dr_gen 放到 demo/ 下。"
    )


VIDEO_DR_SOURCE_ROOT = _resolve_video_dr_source_root()

if str(VIDEO_DR_SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(VIDEO_DR_SOURCE_ROOT))

import assemble_sft_dataset as _video_dr_assemble  # type: ignore  # noqa: E402
import config as _video_dr_config  # type: ignore  # noqa: E402
import utils as _video_dr_utils  # type: ignore  # noqa: E402


VIDEO_DR_SYSTEM_PROMPT_STEP14 = _video_dr_assemble.NEW_SYSTEM_PROMPT.replace(
    "You have a MAXIMUM OF 5 ATTEMPTS (loops) to find the answer.",
    "You have a MAXIMUM OF 10 ATTEMPTS (loops) to find the answer.",
)

VIDEO_DR_STOP_SEARCH_RULES = """

# Stop Searching Rules (STRICT)
- If a search result already contains the requested fact or enough evidence to infer the answer, stop using tools and provide the final `<answer>`.
- Do not repeatedly refine the same web search query after it has already returned the same fact. A near-duplicate query is not useful evidence.
- Near the final attempt, you MUST use the evidence already collected and answer; do not call another tool just to confirm the same fact again.
- If evidence is incomplete but the attempt budget is nearly exhausted, give the best supported answer in `<answer>` and mention uncertainty only inside `<think>`.
"""

VIDEO_DR_SYSTEM_PROMPT_QWEN235B_REPAIR = VIDEO_DR_SYSTEM_PROMPT_STEP14 + VIDEO_DR_STOP_SEARCH_RULES
VIDEO_DR_SYSTEM_PROMPT = VIDEO_DR_SYSTEM_PROMPT_QWEN235B_REPAIR


def get_video_dr_system_prompt(profile: str = "current") -> str:
    """按评测兼容 profile 返回 VideoDR system prompt。"""
    if profile == "step14_plus_tavily432":
        return VIDEO_DR_SYSTEM_PROMPT_STEP14
    return VIDEO_DR_SYSTEM_PROMPT_QWEN235B_REPAIR

MARS_RETRIEVAL_ADDRESS = getattr(_video_dr_config, "MARS_RETRIEVAL_ADDRESS", "")
MARS_SUMMARIZER_ADDRESS = getattr(_video_dr_config, "MARS_SUMMARIZER_ADDRESS", "")
MARS_SUMMARIZER_MODEL = getattr(_video_dr_config, "MARS_SUMMARIZER_MODEL", "")
IMAGE_SEARCH_MODE = getattr(_video_dr_config, "IMAGE_SEARCH_MODE", "")
GATEWAY_URL = getattr(_video_dr_config, "GATEWAY_URL", "")
GATEWAY_USERNAME = getattr(_video_dr_config, "GATEWAY_USERNAME", "")
GATEWAY_USERID = getattr(_video_dr_config, "GATEWAY_USERID", "")
GATEWAY_TOKEN = getattr(_video_dr_config, "GATEWAY_TOKEN", "")

DEFAULT_VIDEO_MAX_RESOLUTION = getattr(_video_dr_config, "DEFAULT_MAX_RESOLUTION", 768)
DEFAULT_VIDEO_JPEG_QUALITY = getattr(_video_dr_config, "DEFAULT_JPEG_QUALITY", 85)
DEFAULT_VIDEO_INITIAL_FRAMES = getattr(_video_dr_config, "DEFAULT_MAX_FRAMES", 64)
DEFAULT_VIDEO_INTERVAL_SAMPLES = getattr(_video_dr_config, "DEFAULT_INTERVAL_SAMPLES", 8)

normalize_bbox = _video_dr_utils.normalize_bbox
get_bbox_config = _video_dr_utils.get_bbox_config
add_search_padding = _video_dr_utils.add_search_padding
crop_frame = _video_dr_utils.crop_frame
real_image_search = _video_dr_utils.real_image_search
mars_web_search = _video_dr_utils.mars_web_search
ImageSearchCache = _video_dr_utils.ImageSearchCache


def extract_video_frames_1fps(
    video_path: str,
    output_dir: str,
    max_resolution: int = DEFAULT_VIDEO_MAX_RESOLUTION,
    jpeg_quality: int = DEFAULT_VIDEO_JPEG_QUALITY,
) -> Dict[int, str]:
    """按 1fps 抽取视频帧，返回 `{frame_index: path}`。"""
    os.makedirs(output_dir, exist_ok=True)

    existing = {}
    for name in sorted(os.listdir(output_dir)):
        if name.startswith("frame_") and name.endswith(".jpg"):
            idx = int(name.replace("frame_", "").replace(".jpg", ""))
            existing[idx] = os.path.join(output_dir, name)
    if existing:
        return existing

    vf = (
        f"fps=1,"
        f"scale='min({max_resolution},iw)':'min({max_resolution},ih)'"
        f":force_original_aspect_ratio=decrease"
    )
    q = max(2, min(31, round(2 + (100 - jpeg_quality) * 29 / 99)))
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        video_path,
        "-vf",
        vf,
        "-q:v",
        str(q),
        "-start_number",
        "0",
        os.path.join(output_dir, "frame_%06d.jpg"),
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
    if proc.returncode != 0:
        raise RuntimeError(f"ffmpeg 1fps 抽帧失败: {proc.stderr[-500:]}")

    frames = {}
    for name in sorted(os.listdir(output_dir)):
        if name.startswith("frame_") and name.endswith(".jpg"):
            idx = int(name.replace("frame_", "").replace(".jpg", ""))
            frames[idx] = os.path.join(output_dir, name)
    if not frames:
        raise RuntimeError(f"未从视频中抽取到任何帧: {video_path}")
    return frames


def uniform_sample_indices(total_frames: int, num_samples: int) -> List[int]:
    """在 `[0, total_frames)` 上均匀采样。"""
    if total_frames <= num_samples:
        return list(range(total_frames))
    return sorted(set(int(i) for i in np.linspace(0, total_frames - 1, num_samples)))


def sample_interval(
    all_frames: Dict[int, str],
    start: int,
    end: int,
    num_samples: int = DEFAULT_VIDEO_INTERVAL_SAMPLES,
) -> List[Tuple[int, str]]:
    """在区间 `[start, end]` 内均匀采样若干帧。"""
    available = sorted(k for k in all_frames if start <= k <= end)
    if not available:
        return []
    if len(available) <= num_samples:
        return [(k, all_frames[k]) for k in available]
    positions = np.linspace(0, len(available) - 1, num_samples, dtype=int)
    selected = sorted(set(available[p] for p in positions))
    return [(k, all_frames[k]) for k in selected]


def get_frame(all_frames: Dict[int, str], idx: int) -> Tuple[int, str]:
    """返回精确帧；若不存在则返回最近邻帧。"""
    if idx in all_frames:
        return idx, all_frames[idx]
    nearest = min(all_frames.keys(), key=lambda key: abs(key - idx))
    return nearest, all_frames[nearest]


def load_videodr_csv_samples(
    annotation_path: str,
    video_root: str,
    dataset_name: str,
) -> List[dict]:
    """读取无表头的 VideoDR CSV，并映射到 `video/<id>.mp4`。"""
    samples: List[dict] = []
    with open(annotation_path, "r", encoding="utf-8-sig", newline="") as f:
        reader = csv.reader(f)
        for row_idx, row in enumerate(reader, start=1):
            if not row:
                continue
            if len(row) < 5:
                raise ValueError(
                    f"VideoDR CSV 第 {row_idx} 行字段不足 5 个: {row}"
                )

            sample_id, question, answer, category, difficulty = [item.strip() for item in row[:5]]
            if not sample_id:
                raise ValueError(f"VideoDR CSV 第 {row_idx} 行缺少样本 id。")

            video_path = os.path.join(video_root, f"{sample_id}.mp4")
            if not os.path.exists(video_path):
                raise FileNotFoundError(
                    f"VideoDR 视频不存在: id={sample_id}, path={video_path}"
                )

            samples.append({
                "id": f"{dataset_name}-{sample_id}",
                "source_id": sample_id,
                "question": question,
                "answer": [answer],
                "dataset": dataset_name,
                "task_kind": "video_dr",
                "video_path": video_path,
                "video_root": video_root,
                "image_path": "",
                "judge_image_path": "",
                "category": category,
                "difficulty": difficulty,
            })

    return samples