model/run.py

# run.py
# -*- coding: utf-8 -*-
import os
import cv2
import csv
import json
import base64
import time
import argparse
from pathlib import Path
from typing import Dict, List, Tuple, Any, Optional

import requests
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed

from parser import parse_step5
from video_utils import draw_inspection_video, draw_mask_video

DEFAULT_SYSTEM_PROMPT = """你是“视频编辑定位规划器（Video Editing Box Planner）”。你的目标是：在多关键帧与编辑指令的条件下，通过链式思维（CoT）严格按 Step1→Step5 推理，产出每个任务唯一且稳定、准确、物理一致的 逐帧 bounding box 序列（风格化任务输出空数组）。
重要总则：
这些关键帧来自同一段视频的时序采样（如等间隔/1fps），输入顺序=时间先后顺序。推理时必须将其视作连续的时间序列，不仅考虑单帧定位，还要考虑前后帧位置/尺度/可见性随时间的合理变化。
仅在 Step5 输出坐标；Step2 只进行分析，不得输出任何坐标。
每个任务每一帧只允许一个 boxing。
多任务必须先拆分，分别推理；跨关键帧保持 ID 一致，不得越界、不得产生不必要重叠。
风格化任务不需要坐标，Step5 对应位置输出 []。
输出语言使用中文。
严格保证 Step5 的最终单行输出格式（见下），Step5 不得附加任何额外文本。
思维链步骤要求（必须逐步输出 Step1–Step4 的分析；坐标仅能出现在 Step5）
Step1｜任务解析（Task Parsing）
判断是单任务还是多任务。
对每个任务做类型识别 ∈ {增加、删除、替换或修改、风格化}。
为每个任务提取受影响对象（语义类别/显著特征）与是否需要坐标（风格化=否，其它=是）。
产出任务清单（按“任务一/任务二/…”编号）：{任务名, 类型, 受影响对象(摘要), 需坐标: 是|否}。
Step2｜图像理解与任务分析（Perception & Task Analysis）——只分析，不给坐标
开展跨关键帧证据融合与任务相关要素分析，不得输出任何坐标。包含四个子步骤：
2.1 时序设定与跨帧对齐
明确：这些关键帧是同一段视频的时序采样帧，输入顺序即时间顺序；
估计全局相机运动（平移/旋转/变焦等），并在时间轴上建立连贯的运动解释；
对候选对象进行实例对齐与 ID 保持；
标注形变属性：rigidity ∈ {rigid, non-rigid}（如人体/布料/液体为 non-rigid）。
2.2 任务相关元素提取
目标对象、参照物（地面/桌面/墙面/容器/人体部位/标牌等）、潜在遮挡体；
区分相似目标的辨识线索：颜色/纹理/轮廓/材质高光/相对尺度/典型位置。
2.3 相对关系与时序轨迹初稿（文字层）
空间关系：上/下/左/右/前/后/接触/包含/平行/垂直；
尺度与透视：与参照物比例、变焦导致的等比缩放 vs 真实位移；
时序连续性与轨迹假设：给出帧间位移方向、速度/加速度趋势、尺度随时间的变化趋势与可见性时间轴（visibility ∈ {visible, partial, hidden}，含遮挡起止）；
要求：除非对象静止或相机完全静止且任务不引入运动，否则不同帧的定位不应完全一致；需有与时序一致的合理位移/缩放解释。
2.4 按任务类型的分析分支（仍不输出坐标）
增加：
在时序上下文中，以构图/遮挡风险/支撑可行/指令一致/相机一致五维度，对候选落点与尺度的时间演化给出文字化理由与排序；
仅在 Step5 选择Top-1 作为最终唯一 boxing，并在所有帧输出连贯的轨迹序列。
删除：
指出应删除实例的辨识要点与反混淆清单；
标注删除后背景修复敏感区；
给出时序上的一致性依据（同一对象在前后帧的可见性/位置演变）。
替换或修改：
说明原物体定位要点与替换后合理尺度/位置在时间轴上的连续变化（保持关键相对关系，或注明允许的微调范围）；
何时需要按时间递进的缩放/微移以保持支撑或避免穿插。
风格化：
仅分析风格域/保护区域/时序一致性/光照一致性；
明确“不需要坐标；Step5 输出 []”。
Step3｜物理与拍摄语法一致性（Physics & Cinematography Consistency）
为每个任务列出相关的现实约束及其对时间维度的影响：
物理（随时间）：重力/支撑接触的持续性，摩擦/滑动的速度变化，动量/碰撞导致的位移与加速度变化，液体水平的稳定性，布料/软体的滞后形变，刚体尺度的连续性；
拍摄语法（随时间）：透视/地平线跨帧一致；变焦→按时间等比例缩放；运动模糊的方向与幅度与观察到的帧间运动一致；阴影/高光与光源方向随时间一致；
反事实自检：若忽略这些时序约束会出现的错误（轨迹跳变、漂浮/穿插、比例突变、阴影方向跳变、背景连接断裂），并给出规避策略（用于约束 Step5 的逐帧 bbox）。
Step4｜指令增强（Prompt Enrichment，禁止包含逐帧/具体帧号信息）
在不改变语义的前提下，生成增强版编辑指令。此处必须满足：
不得出现任何逐帧描述、具体帧号/时间戳、或“在第k帧/在t=…”之类的信息；
仅使用与时间无关的全局性/趋势型表述来帮助编辑模型理解与约束，例如：
相对位置/相对大小的一般范围/趋势（如“相对参照物保持前侧并逐步接近”、“大小与参照物比例保持在X–Y区间”）；
接触/支撑/包含关系与遮挡顺序的一般性表述（如“与地面保持接触，可能被路障部分遮挡”），不指名具体帧；
关键事件的泛化描述（如“出现/消失/接触/遮挡会发生”，但不给出事件对应的帧编号或具体时刻）；
相机一致性：若 camera_hint=zoom，只说明等比例缩放的一般约束，不绑定到具体帧；
不改变区域/保护清单的全局性说明（如“人脸与文字保持不变”）。
目的：让增强指令在任何采样方案与任意帧对齐下都成立，避免把增强指令与具体采样帧强绑定。
Step5｜固定输出格式（唯一可输出坐标的步骤，支持逐帧序列）
只输出一行，必须严格使用以下格式（中文）：
修改后的编辑指令：{增强后的指令文本}，任务一boxing序列：[(t={t0},[{x1,y1,x2,y2}]);(t={t1},[{x1,y1,x2,y2}]);…]，任务二boxing序列：[(t={t0},[{x1,y1,x2,y2}]);(t={t1},[{x1,y1,x2,y2}]);…]，任务三boxing序列：[(t={t0},[{x1,y1,x2,y2}]);…]
关于多帧与坐标的强制规范：
帧顺序：每个任务的 boxing序列 中，帧项必须严格按照输入 keyframes 的顺序列出；{t0},{t1},… 必须与输入中对应的 t 一致。
一帧一个框：每个任务在每个列出的帧上必须且只能有一个 bbox。
时序合理性：除非在 Step2/Step3 已明确论证“静止/无相对运动”，否则不同帧的 bbox 不能完全相同；应体现与任务与相机解释一致的位移/缩放/可见性变化。
遮挡也要输出：即便该帧 visibility='hidden'，也必须输出时序外推的 bbox（不允许缺帧）。
坐标格式：
使用 四元组 [x1,y1,x2,y2]；
0–1 归一化小数（包含 0 与 1），相对于该帧原始图像尺寸；
x1,x2 表示距左边界的相对距离，y1,y2 表示距上边界的相对距离；
0 ≤ x1 < x2 ≤ 1，0 ≤ y1 < y2 ≤ 1。
单/多任务：若仅有单任务，也必须写成“任务一boxing序列：[…]”。
仅此一行：不得在这一行之外追加任何文本或解释。"""

def uniform_sample_indices(total: int, n: int) -> List[int]:
    if total <= 0 or n <= 0:
        return []
    if n >= total:
        return list(range(total))
    step = total / float(n)
    idx = [int(round(i * step)) for i in range(n)]
    idx[-1] = min(idx[-1], total - 1)
    return sorted(set(idx))

def fps1_indices(total: int, fps_src: float) -> List[int]:
    if total <= 0 or fps_src <= 0:
        return []
    return [int(round(i * fps_src)) for i in range(int(total / fps_src) + 1) if int(round(i * fps_src)) < total]

def grab_frames_b64(video_path: str, indices: List[int]) -> Tuple[List[str], float, Tuple[int, int]]:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"视频打不开：{video_path}")
    W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0

    wanted = set(indices)
    b64s = []
    cur = 0
    ok = True
    while ok and wanted:
        ok, frame = cap.read()
        if not ok:
            break
        if cur in wanted:
            _, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
            b64s.append(base64.b64encode(buf.tobytes()).decode("utf-8"))
            wanted.remove(cur)
        cur += 1
    cap.release()
    return b64s, fps, (W, H)

def build_messages(b64_list: List[str], edit_instruction: str) -> List[Dict[str, Any]]:
    content = []
    for b in b64_list:
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{b}"}
        })
    content.append({"type": "text", "text": edit_instruction})
    return content

def call_vllm_chat(api_base: str, model: str, system_prompt: str, messages_user,
                   temperature: float = 0.2, max_tokens: int = 4096, timeout: int = 1200,
                   session: Optional[requests.Session] = None) -> str:
    """
    防截断：客户端上调 max_tokens=4096，HTTP 超时 1200s。
    注：服务端 run.sh 也已设置 --max-model-len=40960、降低 --max-num-seqs，给单请求更多 KV 空间。
    """
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
            {"role": "user",   "content": messages_user}
        ],
        "temperature": temperature,
        "max_tokens": max_tokens
    }
    req = session.post if session else requests.post
    r = req(f"{api_base}/chat/completions",
            headers={"Content-Type": "application/json"},
            json=payload, timeout=timeout)
    r.raise_for_status()
    data = r.json()
    return data["choices"][0]["message"]["content"]

def ensure_dirs(base: Path) -> Dict[str, Path]:
    paths = {
        "inspect": base / "inspection_videos",
        "mask": base / "mask_videos",
        "raw": base / "raw_text",
        "enh": base / "enhanced_text",
        "json": base / "parsed_json",
        "logs": base / "logs"
    }
    for p in paths.values():
        p.mkdir(parents=True, exist_ok=True)
    return paths

def read_input_csv(csv_path: str) -> List[Dict[str, str]]:
    rows = []
    with open(csv_path, "r", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        for r in reader:
            rows.append(r)
    return rows

def get_instruction(row: Dict[str, str]) -> str:
    en = (row.get("instruction") or "").strip()
    zh = (row.get("chinese_instruction") or "").strip()
    if en or zh:
        merged = ",".join([x for x in (en, zh) if x]).strip(",")
        return merged if merged else (row.get("edit_instruction_src") or "").strip()
    return (row.get("edit_instruction_src") or "").strip()

def process_one(row: Dict[str, str], args, paths, session: requests.Session) -> Dict[str, Any]:
    uid = str(row.get("videoid") or row.get("uid") or "").strip()
    vpath = (row.get("source_video_path") or row.get("video_path") or "").strip()
    insrc = get_instruction(row)
    if not uid:
        uid = Path(vpath).stem
    video_name = Path(vpath).stem

    sampling_mode = "fps_1" if args.sampling == "fps_1" else "uniform_5"
    key = f"{uid}::{sampling_mode}"

    result = {
        "key": key,
        "uid": uid,
        "video_path": vpath,
        "edit_instruction_src": insrc,
        "enhanced_instruction": "",
        "sampling_mode": sampling_mode,
        "n_sampled": None,
        "indices_json": "",
        "mapping_json": "",
        "tasks_json": "",
        "latency_s": None,
        "out_inspect": "",
        "out_mask": "",
        "out_json": "",
        "raw_text_path": "",
        "enhanced_text_path": "",
        "parsed_ok": False,
        "error": ""
    }

    if not os.path.exists(vpath):
        result["error"] = f"视频不存在或未提供：{vpath}"
        return result

    # 采样
    cap = cv2.VideoCapture(vpath)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps_src = cap.get(cv2.CAP_PROP_FPS) or 25.0
    cap.release()

    if sampling_mode == "fps_1":
        sampled = fps1_indices(total, fps_src)
    else:
        sampled = uniform_sample_indices(total, 5)

    n = len(sampled)
    mapping = {str(i): sampled[i] for i in range(n)}
    t_list = list(range(n))

    result["n_sampled"] = n
    result["indices_json"] = json.dumps(sampled, ensure_ascii=False)
    result["mapping_json"] = json.dumps(mapping, ensure_ascii=False)

    # 抓帧
    try:
        b64s, fps_used, (W, H) = grab_frames_b64(vpath, sampled)
    except Exception as e:
        result["error"] = f"帧采样失败：{e}"
        return result

    msg_user = build_messages(b64s, insrc)

    # 调模型
    raw_text = ""
    t0 = time.time()
    if args.dry_run:
        raw_text = "DRY-RUN: 未调用模型。"
    else:
        try:
            raw_text = call_vllm_chat(args.api_base, args.model, DEFAULT_SYSTEM_PROMPT, msg_user,
                                      max_tokens=args.max_tokens, timeout=args.http_timeout, session=session)
        except Exception as e:
            result["error"] = str(e)
            return result
    t1 = time.time()
    result["latency_s"] = t1 - t0

    # 保存原始文本
    raw_path = paths["raw"] / f"{uid}__{sampling_mode}.txt"
    with open(raw_path, "w", encoding="utf-8") as f:
        f.write(raw_text)
    result["raw_text_path"] = str(raw_path)

    # 解析 Step5（最后一行）
    lines = [ln.strip() for ln in raw_text.splitlines() if ln.strip()]
    step5 = lines[-1] if lines else ""
    try:
        enhanced, tasks_json = parse_step5(step5, t_list)
    except Exception as e:
        result["error"] = f"解析失败：{e}"
        return result

    # 写增强后的指令（文本与 CSV 字段）
    enh_path = paths["enh"] / f"{video_name}__{uid}__{sampling_mode}.txt"
    with open(enh_path, "w", encoding="utf-8") as f:
        f.write((enhanced or "").strip() + "\n")
    result["enhanced_instruction"] = enhanced
    result["enhanced_text_path"] = str(enh_path)

    # 写解析 JSON
    parsed_payload = {
        "enhanced_instruction": enhanced,
        "tasks": tasks_json,
        "t_list": t_list
    }
    json_path = paths["json"] / f"{uid}__{sampling_mode}.json"
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(parsed_payload, f, ensure_ascii=False, indent=2)
    result["out_json"] = str(json_path)
    result["tasks_json"] = json.dumps(tasks_json, ensure_ascii=False)

    # 生成两路视频
    try:
        inspect_path = paths["inspect"] / f"{video_name}__{uid}__{sampling_mode}__check.mp4"
        mask_path    = paths["mask"]    / f"{video_name}__{uid}__{sampling_mode}__bbox.mp4"
        draw_inspection_video(vpath, str(inspect_path), sampled_idx=sampled, tasks_json=tasks_json, thickness=4)
        draw_mask_video(vpath, str(mask_path), sampled_idx=sampled, tasks_json=tasks_json)
        result["out_inspect"] = str(inspect_path)
        result["out_mask"] = str(mask_path)
        result["parsed_ok"] = True
    except Exception as e:
        result["error"] = f"视频生成失败：{e}"

    return result

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", required=True)
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--api-base", required=True)
    ap.add_argument("--model", required=True, help="必须与 --served-model-name 一致，如 qwen_vl")
    ap.add_argument("--results-csv", required=True)
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--sampling", default="uniform_5", choices=["uniform_5", "fps_1"])
    # 新增：并发与防截断参数
    ap.add_argument("--workers", type=int, default=int(os.environ.get("CLIENT_HTTP_WORKERS", "2")),
                    help="客户端并发线程数（HTTP 并发）。建议 2~4。")
    ap.add_argument("--http-timeout", type=int, default=int(os.environ.get("CLIENT_HTTP_TIMEOUT", "1200")),
                    help="单次请求超时（秒）。默认 1200。")
    ap.add_argument("--max-tokens", type=int, default=int(os.environ.get("CLIENT_MAX_TOKENS", "4096")),
                    help="生成上限 tokens。默认 4096。")
    args = ap.parse_args()

    outdir = Path(args.outdir)
    paths = ensure_dirs(outdir)

    # 断点续做：读取已完成 key
    done_keys = set()
    if args.resume and os.path.exists(args.results_csv):
        with open(args.results_csv, "r", encoding="utf-8") as f:
            try:
                rr = csv.DictReader(f)
                for r in rr:
                    if r.get("parsed_ok") == "True":
                        done_keys.add(r.get("key"))
            except Exception:
                pass

    fieldnames = [
        "key","uid","video_path","edit_instruction_src","enhanced_instruction",
        "sampling_mode","n_sampled","indices_json","mapping_json","tasks_json",
        "latency_s","out_inspect","out_mask","out_json","raw_text_path",
        "enhanced_text_path","parsed_ok","error"
    ]
    write_header = not os.path.exists(args.results_csv)
    out_csv_f = open(args.results_csv, "a", encoding="utf-8", newline="")
    w = csv.DictWriter(out_csv_f, fieldnames=fieldnames)
    if write_header:
        w.writeheader()

    rows = read_input_csv(args.csv)

    # 预过滤（断点续做跳过）
    todo: List[Dict[str, str]] = []
    for r in rows:
        uid = str(r.get("videoid") or r.get("uid") or "").strip() or Path((r.get("source_video_path") or r.get("video_path") or "")).stem
        sampling_mode = "fps_1" if args.sampling == "fps_1" else "uniform_5"
        key = f"{uid}::{sampling_mode}"
        if key in done_keys:
            # 仍然把占位行刷入（保持行为一致）
            w.writerow({
                "key": key, "uid": uid, "video_path": (r.get("source_video_path") or r.get("video_path") or "").strip(),
                "edit_instruction_src": get_instruction(r), "enhanced_instruction": "",
                "sampling_mode": sampling_mode, "n_sampled": None, "indices_json": "",
                "mapping_json": "", "tasks_json": "", "latency_s": None,
                "out_inspect": "", "out_mask": "", "out_json": "",
                "raw_text_path": "", "enhanced_text_path": "",
                "parsed_ok": False, "error": ""
            }); out_csv_f.flush()
            continue
        todo.append(r)

    # HTTP 复用 Session（减少 TLS/握手与粘滞）
    session = requests.Session()

    # 并行跑
    results: List[Dict[str, Any]] = []
    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as ex:
        fut2idx = {ex.submit(process_one, r, args, paths, session): i for i, r in enumerate(todo)}
        for fut in tqdm(as_completed(fut2idx), total=len(fut2idx), desc="Processing"):
            res = fut.result()
            results.append(res)
            w.writerow(res); out_csv_f.flush()

    out_csv_f.close()
    print(f"结果明细写入：{args.results_csv}")

if __name__ == "__main__":
    main()