Proact-VL-Washing/mage_vl_demo/minicpmo_proactive_colab.py

"""
MiniCPM-o 4.5  主动式流媒体视频分析 — Google Colab 可执行脚本
用法：复制到 Colab 笔记本，依次运行各段代码
"""
# ================================================================
# 单元格 1：安装依赖
# ================================================================
import subprocess, sys, os, time, shutil
import numpy as np
import torch

for pkg in [
    "transformers==4.51.0",
    "accelerate",
    "torch>=2.3.0,<=2.8.0",
    "torchaudio<=2.8.0",
    "minicpmo-utils[all]>=1.0.5",
    "librosa", "soundfile", "opencv-python",
]:
    mod = pkg.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].replace("-", "_")
    try:
        __import__(mod)
    except ImportError:
        print(f"[install] {pkg}")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pkg])
print("[✓] 依赖就绪")

# ================================================================
# 单元格 2：上传视频或提供 URL
# ================================================================
from google.colab import files, drive

VIDEO_PATH = None
try:
    drive.mount('/content/drive', timeout_ms=5000)
    _dp = '/content/drive/MyDrive/video.mp4'
    if os.path.exists(_dp):
        VIDEO_PATH = _dp
except Exception:
    pass

if VIDEO_PATH is None:
    print("请上传 mp4 文件：")
    uploaded = files.upload()
    if uploaded:
        VIDEO_PATH = list(uploaded.keys())[0]
        print(f"[✓] 已上传: {VIDEO_PATH}")

if VIDEO_PATH is None:
    URL = input("或粘贴视频 URL: ").strip().strip('"').strip("'")
    if URL:
        VIDEO_PATH = "/content/uploaded.mp4"
        subprocess.run(["wget", "-q", "-O", VIDEO_PATH, URL], check=True)
        print(f"[✓] 已下载: {VIDEO_PATH}")

assert VIDEO_PATH and os.path.exists(VIDEO_PATH), "未找到视频文件，请重试"
print(f"[video] {VIDEO_PATH}  ({os.path.getsize(VIDEO_PATH)/1e6:.1f} MB)")

# ================================================================
# 单元格 3：加载 MiniCPM-o 4.5 模型
# ================================================================
print("[model] 加载中 (首次约 2-3 分钟) ...")
t0 = time.time()
from transformers import AutoModel

MODEL_ID = "openbmb/MiniCPM-o-4_5"
model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    attn_implementation="sdpa",
    torch_dtype=torch.bfloat16,
    init_vision=True,
    init_audio=True,
    init_tts=True,
)
model.eval().cuda()
model.init_tts()
print(f"[✓] 模型加载完成  ({time.time()-t0:.1f}s)")

# ================================================================
# 单元格 4：配置场景 Prompt
# ================================================================
SYSTEM_PROMPT = """\
你是一位专业的衣物护理顾问，正在观看一段洗衣机投放衣物的视频。
你的职责：
1. 仔细观察画面，识别每次被投放的衣物类型（颜色、材质、款式）
2. 每当检测到新衣物被投入时，立即主动输出一条简短的洗涤/护理建议
3. 每条建议以 [MM:SS] 开头标注时间戳
4. 如果没有新衣物出现，请保持沉默，不要输出任何内容
5. 推荐要具体实用，例如："这件白色棉质T恤建议冷水轻柔洗涤，避免漂白"
请用中文回答。"""

# 其他场景示例（取消注释替换即可）：
# SPORTS = """观看体育比赛直播，在每次进球或关键动作时主动给出解说，以 [MM:SS] 开头。"""
# SECURITY = """观看监控视频，发现异常行为（人员闯入、物品移动）时主动告警，以 [MM:SS] 开头。"""
# COOKING = """观看烹饪视频，每完成一个步骤时主动播报进度和下一步建议，以 [MM:SS] 开头。"""

SEGMENT_SEC  = 8.0   # 每段时长（秒）
GENERATE_AUDIO = False  # True 启用 TTS 语音输出
REF_AUDIO     = None   # 可选：参考音频路径（用于声音克隆）

print(f"[config] segment={SEGMENT_SEC}s  audio={'ON' if GENERATE_AUDIO else 'OFF'}")

# ================================================================
# 单元格 5：视频分段 + 流式推理
# ================================================================
print("[infer] 开始推理 ...")
from minicpmo.utils import get_video_frame_audio_segments
import cv2

def get_duration(path: str) -> float:
    cap = cv2.VideoCapture(path)
    fps  = cap.get(cv2.CAP_PROP_FPS)  or 30.0
    n    = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    cap.release()
    return n / fps

def split_segments(path: str, sec: float) -> list[dict]:
    dur = get_duration(path)
    segs, t, i = [], 0.0, 0
    while t < dur - 0.05:
        end = min(dur, t + sec)
        segs.append({"start": t, "end": end, "idx": i})
        t, i = end, i + 1
    return segs

segments = split_segments(VIDEO_PATH, SEGMENT_SEC)
dur = get_duration(VIDEO_PATH)
print(f"  视频时长 {dur:.1f}s → {len(segments)} 段 (每段 {SEGMENT_SEC}s)")

# 声音克隆（可选）
if REF_AUDIO and os.path.exists(REF_AUDIO):
    ref_audio, _ = librosa.load(REF_AUDIO, sr=16000, mono=True)
    model.init_token2wav_cache(ref_audio)
    print(f"  声音克隆: {REF_AUDIO}")

sys_msg = {"role": "system", "content": [SYSTEM_PROMPT]}
session_id = "proactive_session"
model.reset_session()

results = []
t_total = time.time()

for seg in segments:
    # 提取子片段（避免重复解码）
    seg_path = f"/content/_seg_{seg['idx']:04d}.mp4"
    if not os.path.exists(seg_path):
        subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error",
             "-ss", f"{seg['start']}", "-t", f"{SEGMENT_SEC}",
             "-i", VIDEO_PATH, "-c", "copy",
             "-avoid_negative_ts", "make_zero", seg_path],
            capture_output=True,
        )

    print(f"\n[{seg['idx']:02d}] [{seg['start']:.1f}-{seg['end']:.1f}s] ...", end="", flush=True)
    t_seg = time.time()

    try:
        video_frames, audio_segments, stacked_frames = get_video_frame_audio_segments(
            seg_path, stack_frames=1, use_ffmpeg=True, adjust_audio_length=True
        )
    except Exception as e:
        print(f"  ✗ 帧提取失败: {e}"); continue
    if not video_frames:
        print("  (无帧)"); continue

    # Prefill 系统提示
    model.streaming_prefill(session_id=session_id, msgs=[sys_msg])

    # Prefill 视觉+音频内容
    omni = []
    for fi, frame in enumerate(video_frames):
        omni.append(frame)
        if fi < len(audio_segments):
            omni.append(audio_segments[fi])
        if (stacked_frames is not None and fi < len(stacked_frames)
                and stacked_frames[fi] is not None):
            omni.append(stacked_frames[fi])
    omni.append("请按照系统指令，持续观察视频并主动给出建议。")

    last_idx = len(omni) - 1
    for idx, content in enumerate(omni):
        model.streaming_prefill(
            session_id=session_id,
            msgs=[{"role": "user", "content": [content]}],
            omni_mode=True,
            is_last_chunk=(idx == last_idx),
        )

    # 生成回复
    text_parts, audio_chunks = [], []
    try:
        gen = model.streaming_generate(
            session_id=session_id,
            generate_audio=GENERATE_AUDIO,
            use_tts_template=True,
            enable_thinking=False,
            do_sample=True,
            temperature=0.7,
            max_new_tokens=128,
        )
        for item in gen:
            if GENERATE_AUDIO and isinstance(item, tuple) and len(item) == 2:
                wav_chunk, text_chunk = item
                audio_chunks.append(wav_chunk)
                text_parts.append(text_chunk)
            else:
                text_chunk, _done = item
                text_parts.append(text_chunk)
    except Exception as e:
        print(f"  ✗ 生成错误: {e}"); continue

    text = "".join(text_parts).strip()
    dt = time.time() - t_seg

    if text:
        print(f"  ✓ [{dt:.1f}s] {text[:120]}")
    else:
        print(f"  ○ 沉默 [{dt:.1f}s]")

    results.append({"seg_idx": seg["idx"], "start": seg["start"],
                    "end": seg["end"], "text": text, "dt": dt})

# ================================================================
# 单元格 6：输出汇总
# ================================================================
print("\n" + "=" * 60)
print("  主动式分析结果")
print("=" * 60)
key = [r for r in results if r["text"] and r["text"].strip()]
if not key:
    print("  （模型全程未检测到需主动报告的内容）")
    print("  建议：调整 SYSTEM_PROMPT 或增大 SEGMENT_SEC")
for r in key:
    print(f"  [{r['start']:.1f}s] {r['text']}")

total = sum(r["dt"] for r in results)
print(f"\n总计: {total:.1f}s | 有输出节点: {len(key)}/{len(results)}")