build_construct_video.py

"""
For each selected pair image in OpenVE-3M/selected/, find the corresponding
videos, copy them to construct_video/ with sequential names (001.mp4,
001_orig.mp4, 002.mp4, ...), and write a CSV mapping original names to
renamed ones along with the editing prompt.
"""

import csv
import shutil
from pathlib import Path

BASE_DIR = Path("/home/xinyuy/dataset_processing/OpenVE-3M/OpenVE-3M")
SELECTED_DIR = BASE_DIR / "selected"
VIDEOS_DIR = BASE_DIR / "videos"
CSV_DIR = BASE_DIR / "csv_files"
OUTPUT_DIR = Path("/home/xinyuy/dataset_processing/OpenVE-3M/construct_video")

# Map CSV path prefix → actual subfolder name under VIDEOS_DIR
PREFIX_MAP = {
    "local_add": "local_add",
    "local_remove": "local_remove",
    "global_style": "global_style_new",
    "background_change": "background_change",
    "local_change": "local_change",
    "subject-filtered": "subject-filtered",
}

CSV_FILES = [
    CSV_DIR / "local_add.csv",
    CSV_DIR / "local_remove.csv",
    CSV_DIR / "global_style.csv",
    CSV_DIR / "local_add_in_subject.csv",
]


def resolve_video_path(csv_path: str) -> Path | None:
    parts = csv_path.strip().split("/", 1)
    if len(parts) != 2:
        return None
    folder_key, filename = parts
    actual_folder = PREFIX_MAP.get(folder_key)
    if actual_folder is None:
        return None
    return VIDEOS_DIR / actual_folder / filename


def build_lookup() -> dict[str, dict]:
    """
    Build a lookup: video_stem → {edited_path, original_path, prompt}.
    Index both the edited and original stems so we can find a row by either ID.
    """
    lookup: dict[str, dict] = {}
    for csv_file in CSV_FILES:
        if not csv_file.exists():
            continue
        with open(csv_file, newline="", encoding="utf-8") as f:
            reader = csv.DictReader(f, delimiter=";")
            for row in reader:
                edited_csv = row.get("video", "").strip()
                original_csv = row.get("original_video", "").strip()
                prompt = row.get("prompt", "").strip()
                if not edited_csv or not original_csv:
                    continue
                edited_path = resolve_video_path(edited_csv)
                original_path = resolve_video_path(original_csv)
                if edited_path is None or original_path is None:
                    continue

                entry = {
                    "edited_path": edited_path,
                    "original_path": original_path,
                    "prompt": prompt,
                }
                # Index by both stems so we can look up by either ID
                lookup.setdefault(edited_path.stem, []).append(entry)
                lookup.setdefault(original_path.stem, []).append(entry)

    return lookup


def find_entry(id_a: str, id_b: str, lookup: dict) -> dict | None:
    """Find the CSV entry where one ID is the edited video and the other is original."""
    for entry in lookup.get(id_a, []):
        if entry["edited_path"].stem == id_a and entry["original_path"].stem == id_b:
            return entry
        if entry["edited_path"].stem == id_b and entry["original_path"].stem == id_a:
            return entry
    for entry in lookup.get(id_b, []):
        if entry["edited_path"].stem == id_a and entry["original_path"].stem == id_b:
            return entry
        if entry["edited_path"].stem == id_b and entry["original_path"].stem == id_a:
            return entry
    return None


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    selected_images = sorted(SELECTED_DIR.glob("*.jpg"))
    if not selected_images:
        print("No .jpg files found in selected/")
        return

    print(f"Found {len(selected_images)} selected images. Building CSV lookup...")
    lookup = build_lookup()
    print("Lookup built.")

    results = []
    missing = []
    idx = 1

    for img_path in selected_images:
        stem = img_path.stem
        parts = stem.split("__", 1)
        if len(parts) != 2:
            print(f"  [WARN] Cannot parse pair IDs from: {img_path.name}")
            missing.append({"original_image": img_path.name, "reason": "unparseable name"})
            continue

        id_a, id_b = parts
        entry = find_entry(id_a, id_b, lookup)
        if entry is None:
            print(f"  [WARN] No CSV match for: {img_path.name}")
            missing.append({"original_image": img_path.name, "reason": "no CSV match"})
            continue

        edited_path: Path = entry["edited_path"]
        original_path: Path = entry["original_path"]
        prompt: str = entry["prompt"]

        if not edited_path.exists():
            print(f"  [WARN] Edited video missing: {edited_path}")
            missing.append({"original_image": img_path.name, "reason": f"edited video not found: {edited_path}"})
            continue
        if not original_path.exists():
            print(f"  [WARN] Original video missing: {original_path}")
            missing.append({"original_image": img_path.name, "reason": f"original video not found: {original_path}"})
            continue

        tag = f"{idx:03d}"
        edited_dest = OUTPUT_DIR / f"{tag}.mp4"
        original_dest = OUTPUT_DIR / f"{tag}_orig.mp4"

        shutil.copy2(edited_path, edited_dest)
        shutil.copy2(original_path, original_dest)

        results.append({
            "index": tag,
            "renamed_edited": edited_dest.name,
            "original_edited_name": edited_path.name,
            "renamed_original": original_dest.name,
            "original_source_name": original_path.name,
            "prompt": prompt,
        })

        print(f"  [{tag}] {edited_path.name} → {edited_dest.name}")
        print(f"         {original_path.name} → {original_dest.name}")
        idx += 1

    # Write mapping CSV
    out_csv = OUTPUT_DIR / "index.csv"
    fieldnames = [
        "index",
        "renamed_edited",
        "original_edited_name",
        "renamed_original",
        "original_source_name",
        "prompt",
    ]
    with open(out_csv, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(results)

    print(f"\nDone. Copied {len(results)} pairs to {OUTPUT_DIR}")
    print(f"Index saved to {out_csv}")
    if missing:
        print(f"\nSkipped {len(missing)} images:")
        for m in missing:
            print(f"  {m['original_image']}: {m['reason']}")


if __name__ == "__main__":
    main()