examples/wanvideo/model_training/prepare_instancev_iground.py

#!/usr/bin/env python3
"""
Prepare InstanceV training data from iGround processed JSONL.

Outputs per-line JSON with:
{
  "video": "relative/path/to/clip.mp4",
  "prompt": "caption",
  "instance_prompts": ["phrase1", "phrase2", ...],
  "instance_mask_dirs": [
      {"mask_dir": "/abs/path/to/masks", "instance_id": 0, "num_frames": 49},
      ...
  ]
}
"""

import argparse
import json
import math
import os
from pathlib import Path

import imageio.v2 as imageio
from PIL import Image, ImageDraw
from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare InstanceV data from iGround")
    parser.add_argument(
        "--iground_jsonl",
        type=str,
        default="/data/rczhang/PencilFolder/data/iGround/iGround_train_set_processed.jsonl",
        help="Path to iGround processed JSONL.",
    )
    parser.add_argument(
        "--clips_dir",
        type=str,
        default="/data/rczhang/PencilFolder/data/iGround/Clips/train",
        help="Directory containing iGround clips.",
    )
    parser.add_argument(
        "--mask_root_dir",
        type=str,
        default="/data/rczhang/PencilFolder/data/iGround/InstanceMasks/train",
        help="Root directory to store generated instance masks.",
    )
    parser.add_argument(
        "--output_metadata",
        type=str,
        default="/data/rczhang/PencilFolder/data/iGround/instancev_iground_train.jsonl",
        help="Output metadata JSONL path.",
    )
    parser.add_argument(
        "--dataset_base_path",
        type=str,
        default="/data/rczhang/PencilFolder/data",
        help="Base path used by UnifiedDataset (video paths will be relative to this).",
    )
    parser.add_argument(
        "--min_instances",
        type=int,
        default=1,
        help="Minimum number of instances required.",
    )
    parser.add_argument(
        "--max_instances",
        type=int,
        default=None,
        help="Maximum number of instances to keep (None = keep all).",
    )
    parser.add_argument(
        "--overwrite_masks",
        action="store_true",
        help="Overwrite existing masks for a clip.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Limit number of samples for debugging.",
    )
    return parser.parse_args()


def _safe_relpath(path: str, base_path: str) -> str:
    if not base_path:
        return path
    return os.path.relpath(path, base_path)


def _clamp_bbox(bbox, width: int, height: int):
    if not bbox or len(bbox) != 4:
        return None
    x0, y0, x1, y1 = bbox
    left = max(0, int(math.floor(x0)))
    top = max(0, int(math.floor(y0)))
    right = min(width, int(math.ceil(x1)))
    bottom = min(height, int(math.ceil(y1)))
    if right <= left or bottom <= top:
        return None
    return left, top, right, bottom


def _collect_visible_phrases(phrases, labels_per_frame):
    visible = set()
    for labels in labels_per_frame:
        for label in labels:
            visible.add(label)
    return [p for p in phrases if p in visible]


def _write_masks(
    mask_dir: str,
    phrases,
    labels_per_frame,
    bboxes_per_frame,
    width: int,
    height: int,
    overwrite: bool,
):
    if os.path.isdir(mask_dir) and not overwrite:
        return
    os.makedirs(mask_dir, exist_ok=True)

    phrase_set = set(phrases)
    num_frames = len(bboxes_per_frame)
    for frame_idx in range(num_frames):
        labels = labels_per_frame[frame_idx]
        bboxes = bboxes_per_frame[frame_idx]
        frame_map = {}
        for label, bbox in zip(labels, bboxes):
            if label in phrase_set:
                frame_map[label] = bbox

        for inst_id, phrase in enumerate(phrases):
            mask = Image.new("L", (width, height), 0)
            bbox = frame_map.get(phrase)
            if bbox is not None:
                coords = _clamp_bbox(bbox, width, height)
                if coords is not None:
                    draw = ImageDraw.Draw(mask)
                    draw.rectangle(coords, fill=255)
            mask_path = os.path.join(mask_dir, f"{frame_idx:06d}_No.{inst_id}.png")
            mask.save(mask_path)


def _is_video_readable(video_path: str) -> bool:
    try:
        reader = imageio.get_reader(video_path)
        try:
            reader.get_data(0)
        finally:
            reader.close()
    except Exception:
        return False
    return True


def main():
    args = parse_args()

    Path(args.mask_root_dir).mkdir(parents=True, exist_ok=True)
    Path(os.path.dirname(args.output_metadata)).mkdir(parents=True, exist_ok=True)

    processed = 0
    skipped_missing_video = 0
    skipped_instances = 0
    skipped_unreadable = 0
    wrote = 0

    with open(args.iground_jsonl, "r", encoding="utf-8") as f_in, open(
        args.output_metadata, "w", encoding="utf-8"
    ) as f_out:
        for line in tqdm(f_in, desc="Processing iGround"):
            if args.limit is not None and wrote >= args.limit:
                break
            line = line.strip()
            if not line:
                continue
            processed += 1
            sample = json.loads(line)

            video_id = sample["video_id"]
            clip_id = sample["clip_id"]
            clip_name = f"{video_id}_{clip_id}.mp4"
            clip_path = os.path.join(args.clips_dir, clip_name)
            if not os.path.isfile(clip_path):
                skipped_missing_video += 1
                continue
            if not _is_video_readable(clip_path):
                skipped_unreadable += 1
                continue

            phrases = list(sample.get("phrases", []))
            labels_per_frame = sample.get("labels", [])
            bboxes_per_frame = sample.get("bboxes", [])
            if not phrases or not labels_per_frame or not bboxes_per_frame:
                skipped_instances += 1
                continue

            visible_phrases = _collect_visible_phrases(phrases, labels_per_frame)
            if args.max_instances is not None:
                visible_phrases = visible_phrases[: args.max_instances]

            if len(visible_phrases) < args.min_instances:
                skipped_instances += 1
                continue

            width = int(sample["width"])
            height = int(sample["height"])

            mask_dir = os.path.join(args.mask_root_dir, f"{video_id}_{clip_id}_masks")
            _write_masks(
                mask_dir,
                visible_phrases,
                labels_per_frame,
                bboxes_per_frame,
                width,
                height,
                overwrite=args.overwrite_masks,
            )

            instance_mask_dirs = [
                {
                    "mask_dir": mask_dir,
                    "instance_id": inst_id,
                    "num_frames": len(bboxes_per_frame),
                }
                for inst_id in range(len(visible_phrases))
            ]

            entry = {
                "video": _safe_relpath(clip_path, args.dataset_base_path),
                "prompt": sample.get("caption", ""),
                "instance_prompts": visible_phrases,
                "instance_mask_dirs": instance_mask_dirs,
            }
            f_out.write(json.dumps(entry, ensure_ascii=False) + "\n")
            wrote += 1

    print("Done.")
    print(f"Processed: {processed}")
    print(f"Wrote: {wrote}")
    print(f"Skipped (missing video): {skipped_missing_video}")
    print(f"Skipped (unreadable video): {skipped_unreadable}")
    print(f"Skipped (insufficient instances): {skipped_instances}")


if __name__ == "__main__":
    main()