tmp/video-velocity-model/data_processing/tools/umap_alg.py

#!/usr/bin/env python3
"""
umap.py - 提供 UMAP 降维可视化函数，也可作为命令行工具使用。
"""

from __future__ import annotations

import argparse
import re
import subprocess
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import umap

MODEL_LABELS = ("dinov2", "mae", "siglip2", "flux2_ae")
VIDEO_SUFFIXES = (".avi", ".mp4", ".mov", ".mkv", ".webm")
FRAME_INDEX_PATTERN = re.compile(r"\bn:\s*(\d+)\b")
IS_KEYFRAME_PATTERN = re.compile(r"\biskey:\s*(\d+)\b")
PICT_TYPE_PATTERN = re.compile(r"\btype:\s*([A-Z])\b")


def infer_model_label(input_path: Path) -> str | None:
    lower = str(input_path).lower()
    for label in MODEL_LABELS:
        if label in lower:
            return label
    return None


def load_features(input_path: Path, input_key: str) -> np.ndarray:
    if not input_path.exists():
        raise FileNotFoundError(f"输入文件不存在: {input_path}")

    suffix = input_path.suffix.lower()
    if suffix == ".npy":
        return np.load(input_path)
    if suffix == ".npz":
        with np.load(input_path) as data:
            if input_key not in data:
                available = ", ".join(data.files) or "<empty>"
                raise KeyError(
                    f"'{input_key}' 不在 {input_path} 中，可用键: {available}"
                )
            return data[input_key]
    raise ValueError(f"暂不支持的输入格式: {input_path.suffix}，仅支持 .npy/.npz")


def infer_video_stem(input_path: Path) -> str:
    stem = input_path.stem
    if stem.endswith("_patch_tokens"):
        return stem[: -len("_patch_tokens")]
    return stem


def build_video_candidates(input_path: Path, gop_video_root: Path) -> list[Path]:
    video_stem = infer_video_stem(input_path)
    class_name = input_path.parent.name
    candidates: list[Path] = []

    for base_dir in (gop_video_root / class_name, gop_video_root):
        for suffix in VIDEO_SUFFIXES:
            candidate = base_dir / f"{video_stem}{suffix}"
            if candidate not in candidates:
                candidates.append(candidate)

    return candidates


def resolve_gop_video_path(
    input_path: Path,
    gop_video_path: Path | None,
    gop_video_root: Path | None,
) -> Path | None:
    if gop_video_path is not None:
        if not gop_video_path.exists():
            raise FileNotFoundError(f"GOP 视频不存在: {gop_video_path}")
        return gop_video_path

    if gop_video_root is None:
        return None

    candidates = build_video_candidates(input_path, gop_video_root)
    for candidate in candidates:
        if candidate.exists():
            return candidate

    candidate_text = ", ".join(str(path) for path in candidates)
    raise FileNotFoundError(
        f"无法为 {input_path} 推断 GOP 视频路径；已尝试: {candidate_text}"
    )


def resolve_ffmpeg_executable() -> str:
    try:
        import imageio_ffmpeg
    except ImportError as exc:
        raise RuntimeError(
            "启用 GOP 检测需要安装 imageio-ffmpeg，请使用项目虚拟环境运行该脚本。"
        ) from exc
    return imageio_ffmpeg.get_ffmpeg_exe()


def parse_showinfo_keyframes(log_text: str) -> np.ndarray:
    keyframes: list[int] = []
    for line in log_text.splitlines():
        frame_match = FRAME_INDEX_PATTERN.search(line)
        if frame_match is None:
            continue

        is_keyframe_match = IS_KEYFRAME_PATTERN.search(line)
        pict_type_match = PICT_TYPE_PATTERN.search(line)
        is_keyframe = is_keyframe_match is not None and is_keyframe_match.group(1) == "1"
        is_iframe = pict_type_match is not None and pict_type_match.group(1) == "I"
        if is_keyframe or is_iframe:
            keyframes.append(int(frame_match.group(1)))

    return np.asarray(keyframes, dtype=np.int64)


def sanitize_frame_indices(indices: np.ndarray | list[int], n_frames: int) -> np.ndarray:
    values = np.asarray(indices, dtype=np.int64).reshape(-1)
    if values.size == 0:
        raise ValueError("未检测到任何 GOP 起点帧")

    values = np.unique(values)
    valid_values = values[(values >= 0) & (values < n_frames)]
    if valid_values.size == 0:
        raise ValueError(
            f"检测到的 GOP 起点全部超出当前特征帧范围 [0, {n_frames})"
        )

    return valid_values.astype(np.int32, copy=False)


def detect_gop_boundaries(video_path: Path, n_frames: int) -> np.ndarray:
    ffmpeg_exe = resolve_ffmpeg_executable()
    command = [
        ffmpeg_exe,
        "-hide_banner",
        "-loglevel",
        "info",
        "-i",
        str(video_path),
        "-an",
        "-vf",
        "showinfo",
        "-f",
        "null",
        "-",
    ]
    result = subprocess.run(
        command,
        check=False,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True,
    )
    if result.returncode != 0:
        error_tail = "\n".join(result.stderr.strip().splitlines()[-20:])
        raise RuntimeError(
            f"ffmpeg 无法解析 GOP 起点: {video_path}\n{error_tail}"
        )

    return sanitize_frame_indices(parse_showinfo_keyframes(result.stderr), n_frames)


def project_features(features: np.ndarray, projection: str) -> np.ndarray:
    if features.ndim == 2:
        return features

    if features.ndim == 4:
        if projection == "mean_pool":
            return features.mean(axis=(2, 3))
        if projection == "flatten":
            return features.reshape(features.shape[0], -1)

    raise ValueError(
        f"features 形状 {features.shape} 不受支持；当前仅支持 2D (T, D) 或 4D (T, C, H, W) 输入"
    )


def build_default_title(
    input_path: Path,
    projection: str,
    n_neighbors: int,
    min_dist: float,
) -> str:
    model_label = infer_model_label(input_path)
    prefix = f"{model_label} | " if model_label else ""
    return (
        f"{prefix}{input_path.stem} | {projection} | "
        f"n_neighbors={n_neighbors} | min_dist={min_dist}"
    )


def resolve_output_path(
    input_path: Path,
    output: Path | None,
    output_dir: Path | None,
    projection: str,
    multiple_inputs: bool,
) -> Path:
    if output is not None:
        if multiple_inputs:
            raise ValueError("同时处理多个输入时不能使用 --output，请改用 --output-dir")
        return output

    model_label = infer_model_label(input_path)
    filename = f"{input_path.stem}_{projection}_umap.png"

    if output_dir is None and not multiple_inputs:
        return Path("umap_plot.png")

    base_dir = output_dir if output_dir is not None else Path.cwd()
    if model_label:
        return base_dir / model_label / filename
    return base_dir / filename


def visualize_umap(
    features,
    output_path,
    n_neighbors=15,
    min_dist=0.1,
    metric="euclidean",
    title=None,
    show_colorbar=True,
    cmap="viridis",
    point_size=5,
    alpha=0.7,
    random_state=42,
    gop_boundaries=None,
):
    """
    对二维帧特征进行 UMAP 降维并保存散点图，颜色表示帧顺序。

    参数:
        features: ndarray, shape (n_frames, feature_dim)
        output_path: str | Path, 输出图像路径
        n_neighbors: int, UMAP 邻域大小
        min_dist: float, UMAP 最小距离
        metric: str, 距离度量
        title: str, 图表标题
        show_colorbar: bool, 是否显示颜色条
        cmap: str, 颜色映射
        point_size: float, 点的大小
        alpha: float, 点的透明度
        random_state: int, 随机种子
        gop_boundaries: ndarray | None, GOP 起点帧索引
    """
    if features.ndim != 2:
        raise ValueError("features 应为二维数组 (n_frames, feature_dim)")

    n_frames = features.shape[0]
    if n_frames < 2:
        raise ValueError("UMAP 至少需要 2 帧特征")

    effective_neighbors = min(n_neighbors, max(2, n_frames - 1))
    if effective_neighbors != n_neighbors:
        print(
            f"n_neighbors={n_neighbors} 对 {n_frames} 帧过大，自动调整为 {effective_neighbors}"
        )

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    print(
        f"运行UMAP降维: shape={features.shape}, "
        f"n_neighbors={effective_neighbors}, min_dist={min_dist}, metric={metric}"
    )
    reducer = umap.UMAP(
        n_neighbors=effective_neighbors,
        min_dist=min_dist,
        metric=metric,
        n_components=2,
        random_state=random_state,
    )
    embedding = reducer.fit_transform(features)
    print("降维完成，嵌入形状:", embedding.shape)

    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(
        embedding[:, 0],
        embedding[:, 1],
        c=np.arange(n_frames),
        cmap=cmap,
        s=point_size,
        alpha=alpha,
    )
    if show_colorbar:
        cbar = plt.colorbar(scatter)
        cbar.set_label("frame index")

    if gop_boundaries is not None:
        plt.scatter(
            embedding[gop_boundaries, 0],
            embedding[gop_boundaries, 1],
            s=max(point_size * 12, 36),
            facecolors="none",
            edgecolors="#d62728",
            linewidths=1.5,
            marker="o",
            label="GOP start",
            zorder=3,
        )
        plt.legend(loc="best")

    if title is None:
        title = f"UMAP (n_neighbors={effective_neighbors}, min_dist={min_dist})"
    plt.title(title)
    plt.xlabel("UMAP1")
    plt.ylabel("UMAP2")
    plt.tight_layout()

    plt.savefig(output_path, dpi=150)
    plt.close()
    print(f"图像已保存至: {output_path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="使用 UMAP 可视化视频帧的潜在表示")
    parser.add_argument(
        "--input",
        nargs="+",
        required=True,
        help="输入特征文件路径，支持 .npy 或 .npz，可一次传多个",
    )
    parser.add_argument(
        "--input-key",
        type=str,
        default="features",
        help="读取 .npz 时使用的键名（默认: features）",
    )
    parser.add_argument(
        "--projection",
        type=str,
        choices=("mean_pool", "flatten"),
        default="mean_pool",
        help="当输入为 4D (T, C, H, W) 时，将其投影到二维帧特征的方式",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="单个输入时的输出图像路径；多输入场景请改用 --output-dir",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="多输入时的输出目录；会按模型名自动分子目录，并按输入文件名生成 PNG",
    )
    parser.add_argument(
        "--n_neighbors",
        type=int,
        default=15,
        help="UMAP n_neighbors 参数（默认 15）",
    )
    parser.add_argument(
        "--min_dist",
        type=float,
        default=0.1,
        help="UMAP min_dist 参数（默认 0.1）",
    )
    parser.add_argument(
        "--metric",
        type=str,
        default="euclidean",
        help="UMAP 距离度量（默认 euclidean）",
    )
    parser.add_argument(
        "--no_colorbar",
        action="store_true",
        help="不显示颜色条",
    )
    parser.add_argument(
        "--title",
        type=str,
        default=None,
        help="图表标题；默认自动包含模型名、输入文件名和 UMAP 参数",
    )
    parser.add_argument(
        "--gop-video-path",
        type=Path,
        default=None,
        help="单输入时显式指定原始视频路径；会从视频码流中检测 GOP 起点",
    )
    parser.add_argument(
        "--gop-video-root",
        type=Path,
        default=None,
        help=(
            "按输入特征文件路径自动回推原始视频目录；"
            "例如 <class>/<video>_patch_tokens.npz -> <gop_video_root>/<class>/<video>.*"
        ),
    )
    return parser.parse_args()


def main():
    args = parse_args()
    input_paths = [Path(path).expanduser() for path in args.input]
    multiple_inputs = len(input_paths) > 1
    gop_video_path = args.gop_video_path.expanduser() if args.gop_video_path is not None else None
    gop_video_root = args.gop_video_root.expanduser() if args.gop_video_root is not None else None

    if multiple_inputs and gop_video_path is not None:
        raise ValueError("同时处理多个输入时不能使用 --gop-video-path，请改用 --gop-video-root")

    for input_path in input_paths:
        raw_features = load_features(input_path, args.input_key)
        features = project_features(raw_features, args.projection)
        output_path = resolve_output_path(
            input_path=input_path,
            output=args.output,
            output_dir=args.output_dir,
            projection=args.projection,
            multiple_inputs=multiple_inputs,
        )
        title = args.title or build_default_title(
            input_path=input_path,
            projection=args.projection,
            n_neighbors=args.n_neighbors,
            min_dist=args.min_dist,
        )
        gop_boundaries = None
        video_path = resolve_gop_video_path(
            input_path=input_path,
            gop_video_path=gop_video_path,
            gop_video_root=gop_video_root,
        )
        if video_path is not None:
            print(f"检测 GOP 起点: {video_path}")
            gop_boundaries = detect_gop_boundaries(video_path, features.shape[0])
            print(f"检测到 {len(gop_boundaries)} 个 GOP 起点帧: {gop_boundaries.tolist()}")
        print(
            f"处理 {input_path}: 原始形状={raw_features.shape}, 投影后形状={features.shape}"
        )
        visualize_umap(
            features,
            output_path,
            n_neighbors=args.n_neighbors,
            min_dist=args.min_dist,
            metric=args.metric,
            title=title,
            show_colorbar=not args.no_colorbar,
            gop_boundaries=gop_boundaries,
        )


if __name__ == "__main__":
    main()