examples/wanvideo/model_training/prepare_instancev_data.py

#!/usr/bin/env python3
"""
InstanceV 训练数据预处理脚本

功能：
1. 读取 InstanceCap.jsonl（包含 Video, Global Description, Structural Description）
2. 匹配 InstanceLabel 中的 mask 目录
3. 匹配 OpenVid1M-Video 中的视频文件
4. 生成适合训练的 JSONL 格式

数据目录结构：
  data/InstanceCap/InstanceCap.jsonl          # 原始标注
  data/InstanceLabel/{video_name}_masks/      # 每个 instance 的 mask 序列
  data/OpenVid1M-Video/{video_name}.mp4       # 原始视频

输出格式（每行一个 JSON）：
{
    "video": "path/to/video.mp4",
    "prompt": "global description",
    "instance_prompts": ["instance 0 description", "instance 1 description", ...],
    "instance_mask_dirs": ["path/to/masks/No.0", "path/to/masks/No.1", ...]
}
"""

import os
import json
import argparse
from pathlib import Path
from collections import defaultdict
from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare InstanceV training data")
    parser.add_argument(
        "--instancecap_path",
        type=str,
        default="/data/rczhang/PencilFolder/data/InstanceCap/InstanceCap.jsonl",
        help="Path to InstanceCap.jsonl",
    )
    parser.add_argument(
        "--instance_label_dir",
        type=str,
        default="/data/rczhang/PencilFolder/data/InstanceLabel",
        help="Directory containing instance masks",
    )
    parser.add_argument(
        "--video_dir",
        type=str,
        default="/data/rczhang/PencilFolder/data/OpenVid1M-Video",
        help="Directory containing source videos",
    )
    parser.add_argument(
        "--output_path",
        type=str,
        default="/data/rczhang/PencilFolder/data/instancev_train.jsonl",
        help="Output JSONL path",
    )
    parser.add_argument(
        "--min_instances",
        type=int,
        default=1,
        help="Minimum number of instances required",
    )
    parser.add_argument(
        "--max_instances",
        type=int,
        default=10,
        help="Maximum number of instances to keep",
    )
    parser.add_argument(
        "--use_dense_caption",
        action="store_true",
        help="Use dense caption format for instance prompts",
    )
    return parser.parse_args()


def get_video_name_from_path(video_path: str) -> str:
    """从 video 路径提取 video name（不含扩展名）"""
    return Path(video_path).stem


def find_mask_dirs(instance_label_dir: str, video_name: str) -> dict:
    """
    查找某个视频对应的所有 instance mask 目录
    
    mask 目录结构: {video_name}_masks/
    mask 文件命名: {frame_id:06d}_No.{instance_id}.png
    
    Returns:
        dict: {instance_id: [mask_file_paths]}
    """
    mask_dir = os.path.join(instance_label_dir, f"{video_name}_masks")
    if not os.path.isdir(mask_dir):
        return {}
    
    instance_masks = defaultdict(list)
    for fname in sorted(os.listdir(mask_dir)):
        if not fname.endswith(".png"):
            continue
        # 解析文件名: 000000_No.0.png
        parts = fname.replace(".png", "").split("_No.")
        if len(parts) != 2:
            continue
        frame_id_str, inst_id_str = parts
        try:
            frame_id = int(frame_id_str)
            inst_id = int(inst_id_str)
        except ValueError:
            continue
        instance_masks[inst_id].append(os.path.join(mask_dir, fname))
    
    return dict(instance_masks)


def build_instance_prompt(instance_info: dict) -> str:
    """
    从 InstanceCap 的 instance 信息构建 prompt
    
    instance_info 结构:
    {
        "Class": "person",
        "Appearance": "...",
        "Actions and Motion": "...",
        "Position": "..."
    }
    """
    cls = instance_info.get("Class", "object")
    appearance = instance_info.get("Appearance", "")
    actions = instance_info.get("Actions and Motion", "")
    
    # 构建精简但信息丰富的 prompt
    parts = []
    if cls:
        parts.append(f"A {cls}")
    if appearance:
        parts.append(appearance)
    if actions:
        parts.append(actions)
    
    prompt = ". ".join(parts)
    # 清理多余空格和标点
    prompt = prompt.replace("..", ".").replace("  ", " ").strip()
    if prompt and not prompt.endswith("."):
        prompt += "."
    return prompt


def process_sample(
    sample: dict,
    instance_label_dir: str,
    video_dir: str,
    min_instances: int,
    max_instances: int,
) -> dict | None:
    """
    处理单个样本，返回训练格式的 dict 或 None
    """
    video_name = get_video_name_from_path(sample.get("Video", ""))
    if not video_name:
        return None
    
    # 检查视频是否存在
    video_path = os.path.join(video_dir, f"{video_name}.mp4")
    if not os.path.isfile(video_path):
        return None
    
    # 查找 mask 目录
    instance_masks = find_mask_dirs(instance_label_dir, video_name)
    if len(instance_masks) < min_instances:
        return None
    
    # 提取 instance 信息
    struct_desc = sample.get("Structural Description", {})
    main_instances = struct_desc.get("Main Instance", {})
    
    instance_prompts = []
    instance_mask_dirs = []
    
    for inst_key in sorted(main_instances.keys()):
        # inst_key: "No.0", "No.1", ...
        try:
            inst_id = int(inst_key.replace("No.", ""))
        except ValueError:
            continue
        
        if inst_id not in instance_masks:
            continue
        
        inst_info = main_instances[inst_key]
        prompt = build_instance_prompt(inst_info)
        if not prompt:
            continue
        
        instance_prompts.append(prompt)
        # 存储整个 mask 目录路径（而非单个文件列表，训练时动态加载）
        mask_dir = os.path.dirname(instance_masks[inst_id][0])
        instance_mask_dirs.append({
            "mask_dir": mask_dir,
            "instance_id": inst_id,
            "num_frames": len(instance_masks[inst_id]),
        })
        
        if len(instance_prompts) >= max_instances:
            break
    
    if len(instance_prompts) < min_instances:
        return None
    
    # 构建输出
    global_desc = sample.get("Global Description", "")
    background = struct_desc.get("Background Detail", "")
    camera = struct_desc.get("Camera Movement", "")
    
    # 合并为完整 prompt
    full_prompt_parts = [global_desc]
    if background:
        full_prompt_parts.append(background)
    if camera:
        full_prompt_parts.append(camera)
    full_prompt = " ".join(full_prompt_parts)
    
    return {
        "video": video_path,
        "prompt": full_prompt,
        "instance_prompts": instance_prompts,
        "instance_mask_dirs": instance_mask_dirs,
    }


def main():
    args = parse_args()
    
    # 读取 InstanceCap
    print(f"Loading InstanceCap from {args.instancecap_path}")
    samples = []
    with open(args.instancecap_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                samples.append(json.loads(line))
            except json.JSONDecodeError as e:
                print(f"Warning: Failed to parse line: {e}")
                continue
    
    print(f"Loaded {len(samples)} samples")
    
    # 处理每个样本
    output_samples = []
    for sample in tqdm(samples, desc="Processing samples"):
        result = process_sample(
            sample,
            args.instance_label_dir,
            args.video_dir,
            args.min_instances,
            args.max_instances,
        )
        if result is not None:
            output_samples.append(result)
    
    print(f"Valid samples: {len(output_samples)} / {len(samples)}")
    
    # 写入输出
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    with open(args.output_path, "w", encoding="utf-8") as f:
        for sample in output_samples:
            f.write(json.dumps(sample, ensure_ascii=False) + "\n")
    
    print(f"Saved to {args.output_path}")
    
    # 打印统计信息
    if output_samples:
        avg_instances = sum(len(s["instance_prompts"]) for s in output_samples) / len(output_samples)
        print(f"Average instances per sample: {avg_instances:.2f}")


if __name__ == "__main__":
    main()