build_style_transfer_reference.py

#!/usr/bin/env python3
"""Build `style-transfer-reference` straight from the raw OpenVE-3M global_style CSV.

Same selection pipeline as `regenerate_style_transfer.py` (style extraction, canonical-19
filter, N-per-style sampling with no `original_video` shared across styles), so the
manifest keeps full provenance back to the source content-hash paths.  What changes is the
task the samples encode:

    regenerate_style_transfer.py :  (original video + prompt) -> stylized video
    this script                  :  (style reference image + prompt) -> stylized video

The stylized clip is the generation target (`media_path`).  Its conditioning is a single
frame sampled at random from a *different* stylized clip of the *same* style
(`reference_path`) — drawing the frame from a foreign clip is what keeps the pair honest,
since a model cannot then solve the task by copying content out of its own reference.

The `_org` source clips are not copied (they have no role once conditioning is an image),
but every row still records `original_org` so they can be recovered from the CSV.

`--horizontal-only` restricts sampling to landscape clips (width > height) and writes to
`style-transfer-reference-horizontal` by default.  Reference frames are cut from another
clip of the same style, so filtering the clips is enough to make every reference landscape
too.
"""

from __future__ import annotations

import argparse
import json
import os
import random
import re
import shutil
import sys
import time
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from dataclasses import dataclass

import cv2
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_VIDEO_BASE = "/mnt/data/xinyuy/datasets/OpenVE-3M/videos"
DEFAULT_CSV = "/mnt/data/xinyuy/datasets/OpenVE-3M/csv_files/global_style.csv"
DEFAULT_DST = os.path.join(HERE, "style-transfer-reference")
DEFAULT_DST_HORIZONTAL = os.path.join(HERE, "style-transfer-reference-horizontal")

# ── Style taxonomy (kept identical to regenerate_style_transfer.py) ──────────────
CANONICAL_19 = {
    'Abstract Art', 'Ancient Style', 'Cartoon', 'Chinese Ink Wash Painting',
    'Cubist', 'Cyberpunk', 'Ghibli', 'Gongbi', 'Impressionist', 'Oil Painting',
    'Pixel Art', 'Pointillist', 'Pop Art', 'Sketch', 'Steampunk', 'Surrealist',
    'Thick-Paint Hand-Drawn', 'Ukiyoe', 'Watercolor',
}

_STOP = (r'(?:\s+(?:aesthetic|animation style|animation|style|Style|visual style|visual|look'
         r'|filter|effect|technique|art style|principles|to this)|\s*,|\s*\.)')

_NORMALIZE = {
    # NOTE: regenerate_style_transfer.py maps 'dynamic' -> 'Abstract Art'.  That rule is
    # dropped here: _STOP breaks the capture at the first comma, so "Apply a dynamic,
    # immersive snowy style" yields the bare modifier 'dynamic' and every such prompt —
    # snowy, rainy, sunlit, sketch — was filed as Abstract Art.  Without the rule those
    # rows fail the CANONICAL_19 filter and drop out of sampling.
    'principles of abstract art':  'Abstract Art',
    'Aesthetic Ancient-style':     'Ancient Style',
    'iconic Ghibli':               'Ghibli',
    'Studio Ghibli':               'Ghibli',
    'traditional Gongbi painting': 'Gongbi',
    'sketch':                      'Sketch',
    'Pointillism':                 'Pointillist',
    'Thick-paint Hand-drawn':      'Thick-Paint Hand-Drawn',
    'pixel':                       'Pixel Art',
    'retro pixel':                 'Pixel Art',
    'Pixel':                       'Pixel Art',
    'oil painting':                'Oil Painting',
    'Chinese Ink Wash':            'Chinese Ink Wash Painting',
}

STYLE_PREFIX = {
    'Abstract Art':              'abstract',
    'Ancient Style':             'ancient',
    'Cartoon':                   'cartoon',
    'Chinese Ink Wash Painting': 'chinese_ink_wash',
    'Cubist':                    'cubist',
    'Cyberpunk':                 'cyberpunk',
    'Ghibli':                    'ghibli',
    'Gongbi':                    'gongbi',
    'Impressionist':             'impressionist',
    'Oil Painting':              'oil_painting',
    'Pixel Art':                 'pixel_art',
    'Pointillist':               'pointillist',
    'Pop Art':                   'pop_art',
    'Sketch':                    'sketch',
    'Steampunk':                 'steampunk',
    'Surrealist':                'surrealist',
    'Thick-Paint Hand-Drawn':    'thick_paint',
    'Ukiyoe':                    'ukiyoe',
    'Watercolor':                'watercolor',
}

# A frame that is nearly a flat colour (fade-in/out, black leader) tells a model nothing
# about the style, so resample a few times before giving up on a clip.
FLAT_FRAME_STD = 6.0
FLAT_FRAME_RETRIES = 5


def extract_style(prompt: str):
    m = re.search(r'Apply (?:the |a |an )(.+?)' + _STOP, prompt)
    if m:
        return m.group(1).strip().strip("'\"")
    m2 = re.search(r'Apply (?:the |a |an )((?:\S+ ?){1,4})', prompt)
    if m2:
        return m2.group(1).strip().strip("'\"")
    m3 = re.search(r'Convert this (.+?)(?:-animated|-style\b| video\b)', prompt)
    if m3:
        return m3.group(1).strip().strip("'\"")
    return None


def remap(path: str) -> str:
    return path.replace('global_style/', 'global_style_new/', 1)


def resolve(video_base: str, rel_path: str):
    """CSV paths point at `global_style/`; the clips actually live in `global_style_new/`."""
    full = os.path.join(video_base, remap(rel_path))
    if os.path.exists(full):
        return full
    fallback = os.path.join(video_base, rel_path)
    return fallback if os.path.exists(fallback) else None


# ── Orientation probing (only used by --horizontal-only) ────────────────────────
def probe_dims(video_base: str, rel_path: str):
    """Return `(width, height)` for a clip, or None when it cannot be read."""
    full = resolve(video_base, rel_path)
    if full is None:
        return None
    cap = cv2.VideoCapture(full)
    if not cap.isOpened():
        return None
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    if w <= 0 or h <= 0:  # a few containers only report a size once decoding starts
        ok, frame = cap.read()
        if ok and frame is not None:
            h, w = frame.shape[:2]
    cap.release()
    return (w, h) if w > 0 and h > 0 else None


def is_landscape(video_base: str, rel_path: str) -> bool:
    dims = probe_dims(video_base, rel_path)
    return dims is not None and dims[0] > dims[1]


def sample_landscape(available: pd.DataFrame, need: int, video_base: str,
                     workers: int, seed: int) -> tuple[pd.DataFrame, int]:
    """Draw `need` rows whose stylized clip is landscape, probing lazily.

    Returns the sampled rows and how many clips had to be opened to find them.

    Probing the whole candidate pool up front would mean opening ~100k clips.  Since only a
    small minority are portrait, walking a shuffled order in batches reaches `need` keepers
    after roughly `need` probes per style and leaves the rest of the pool untouched.
    """
    order = available.sample(frac=1, random_state=seed).reset_index(drop=True)
    keep_idx: list[int] = []
    cursor = probed = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        while cursor < len(order) and len(keep_idx) < need:
            short = need - len(keep_idx)
            batch = min(len(order) - cursor, max(short + short // 2 + 16, 32))
            chunk = order.iloc[cursor:cursor + batch]
            flags = list(pool.map(lambda p: is_landscape(video_base, p), chunk['video']))
            for offset, ok in enumerate(flags):
                if ok:
                    keep_idx.append(cursor + offset)
                    if len(keep_idx) == need:
                        break
            cursor += batch
            probed += batch
    return order.iloc[keep_idx], probed


# ── Step 1-3: select samples ────────────────────────────────────────────────────
def select_samples(csv_path: str, video_base: str, per_style: int, sample_seed: int,
                   horizontal_only: bool = False, workers: int = 16) -> pd.DataFrame:
    print("Loading CSV...")
    data = pd.read_csv(csv_path, sep=';')

    def file_exists(rel_path):
        return os.path.exists(os.path.join(video_base, remap(rel_path)))

    mask = data['video'].apply(file_exists) & data['original_video'].apply(file_exists)
    filtered = data[mask].reset_index(drop=True)
    print(f"Rows with both videos present: {len(filtered):,}")

    filtered = filtered.copy()
    filtered['style_raw'] = filtered['prompt'].apply(extract_style)
    filtered['style'] = filtered['style_raw'].apply(lambda r: _NORMALIZE.get(r, r) if r else None)

    gs_common = filtered[
        filtered['style'].isin(CANONICAL_19) &
        filtered['prompt'].str.startswith('Apply')
    ]
    print(f"Forward-only rows across 19 styles: {len(gs_common):,}")

    if horizontal_only:
        print(f"Landscape-only sampling ({workers} probe threads)...")

    # Style order is shuffled so no single style gets first pick of the shared originals.
    used_originals: set = set()
    selected = []
    order = pd.Series(sorted(CANONICAL_19)).sample(frac=1, random_state=sample_seed).tolist()
    for style in order:
        style_df = gs_common[gs_common['style'] == style]
        available = style_df[~style_df['original_video'].isin(used_originals)]
        if horizontal_only:
            sampled, probed = sample_landscape(available, per_style, video_base,
                                               workers, sample_seed)
            print(f"  {style:<30} {len(sampled):>4} landscape / {probed:>5} probed", flush=True)
            if len(sampled) < per_style:
                print(f"WARNING: {style!r} only has {len(sampled)} non-overlapping landscape "
                      f"rows (need {per_style})")
        elif len(available) < per_style:
            print(f"WARNING: {style!r} only has {len(available)} non-overlapping rows "
                  f"(need {per_style})")
            sampled = available.sample(frac=1, random_state=sample_seed)
        else:
            sampled = available.sample(n=per_style, random_state=sample_seed)
        used_originals.update(sampled['original_video'].tolist())
        selected.append(sampled)

    final_df = pd.concat(selected).reset_index(drop=True)
    assert final_df.groupby('original_video')['style'].nunique().max() == 1, "overlap detected!"
    print(f"\nSelected {len(final_df):,} rows across {final_df['style'].nunique()} styles")
    return final_df


# ── Step 4: pair each target with a foreign clip of the same style ──────────────
@dataclass
class Job:
    name: str            # e.g. "abstract_000"
    style: str           # display name, e.g. "Abstract Art"
    index: int
    video: str           # target filename inside the dataset
    src_video: str       # absolute path of the stylized clip in VIDEO_BASE
    original_video: str  # CSV path of the stylized clip (provenance)
    original_org: str    # CSV path of the unstylized source clip (provenance)
    prompt: str
    reference_image: str
    donor_video: str = ""   # dataset-relative clip the reference frame is cut from
    donor_name: str = ""


def plan_jobs(final_df: pd.DataFrame, video_base: str, ext: str, seed: int) -> list[Job]:
    jobs: list[Job] = []
    missing = 0
    for style, grp in final_df.groupby('style'):
        prefix = STYLE_PREFIX[style]
        for idx, (_, row) in enumerate(grp.iterrows()):
            src = resolve(video_base, row['video'])
            if src is None:
                print(f"  MISSING: {row['video']}")
                missing += 1
                continue
            stem = f"{prefix}_{idx:03d}"
            jobs.append(Job(
                name=stem,
                style=style,
                index=idx,
                video=f"{stem}.mp4",
                src_video=src,
                original_video=row['video'],
                original_org=row['original_video'],
                prompt=row['prompt'],
                reference_image=f"{stem}_ref.{ext}",
            ))
    if missing:
        print(f"[warn] dropped {missing} rows whose stylized clip could not be resolved")

    # Donors come from a shuffled rotation rather than independent draws, so every clip is
    # used as a reference exactly once and no clip's frames dominate the distribution.
    by_style: dict[str, list[Job]] = defaultdict(list)
    for j in jobs:
        by_style[j.style].append(j)

    rng = random.Random(seed)
    for style in sorted(by_style):
        group = by_style[style]
        if len(group) < 2:
            print(f"[warn] style {style!r} has <2 clips, cannot pick a foreign reference")
            continue
        order = list(range(len(group)))
        rng.shuffle(order)
        # Rotating the shuffled order by one guarantees donor != target for every entry.
        for pos, i in enumerate(order):
            donor = group[order[(pos + 1) % len(order)]]
            assert donor.name != group[i].name
            group[i].donor_video = donor.video
            group[i].donor_name = donor.name

    jobs.sort(key=lambda j: j.name)
    return jobs


# ── Step 5: copy clips, then cut reference frames ───────────────────────────────
def copy_video(job: Job, dst: str) -> dict:
    out = os.path.join(dst, job.video)
    try:
        if not os.path.exists(out):
            tmp = f"{out}.tmp{os.getpid()}"
            shutil.copy2(job.src_video, tmp)
            os.replace(tmp, out)
        return {"name": job.name, "status": "ok"}
    except Exception as exc:
        return {"name": job.name, "status": "error", "error": f"copy: {type(exc).__name__}: {exc}"}


def _read_frame(cap: cv2.VideoCapture, idx: int):
    """Seek to `idx`; fall back to a sequential scan when the seek lands nowhere."""
    cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
    ok, frame = cap.read()
    if ok and frame is not None:
        return frame
    cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
    last = None
    for _ in range(idx + 1):
        ok, frame = cap.read()
        if not ok or frame is None:
            break
        last = frame
    return last


def extract_reference(job: Job, dst: str, seed: int, jpeg_quality: int,
                      require_landscape: bool = False) -> dict:
    try:
        return _extract_reference(job, dst, seed, jpeg_quality, require_landscape)
    except Exception as exc:  # one bad clip must not abort the whole build
        return {"name": job.name, "status": "error", "error": f"{type(exc).__name__}: {exc}"}


def _extract_reference(job: Job, dst: str, seed: int, jpeg_quality: int,
                       require_landscape: bool) -> dict:
    out_img = os.path.join(dst, job.reference_image)
    if os.path.exists(out_img) and os.path.getsize(out_img) > 0:
        stale = False
        if require_landscape:
            # Under --no-clean a leftover portrait reference from an earlier build would
            # otherwise survive; re-cut it from the (landscape) donor instead.
            cached = cv2.imread(out_img)
            stale = cached is None or cached.shape[1] <= cached.shape[0]
        if not stale:
            return {"name": job.name, "status": "cached", "frame_index": -1, "total_frames": -1}

    cap = cv2.VideoCapture(os.path.join(dst, job.donor_video))
    if not cap.isOpened():
        return {"name": job.name, "status": "error", "error": f"cannot open {job.donor_video}"}

    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1

    # Seed off the sample name so results are identical no matter how work is scheduled.
    rng = random.Random(f"{seed}:{job.name}")
    frame, frame_index = None, 0
    for _ in range(FLAT_FRAME_RETRIES + 1):
        frame_index = rng.randrange(total)
        candidate = _read_frame(cap, frame_index)
        if candidate is None:
            continue
        frame = candidate
        if candidate.std() >= FLAT_FRAME_STD:
            break
    cap.release()

    if frame is None:
        return {"name": job.name, "status": "error",
                "error": f"no decodable frame in {job.donor_video}"}

    if require_landscape and frame.shape[1] <= frame.shape[0]:
        h, w = frame.shape[:2]
        return {"name": job.name, "status": "error",
                "error": f"donor {job.donor_video} is {w}x{h}, not landscape"}

    params = [cv2.IMWRITE_JPEG_QUALITY, jpeg_quality] if out_img.endswith((".jpg", ".jpeg")) else []
    # cv2 picks its encoder from the extension, so the temp name must keep it last.
    base, ext = os.path.splitext(out_img)
    tmp = f"{base}.tmp{os.getpid()}{ext}"
    if not cv2.imwrite(tmp, frame, params):
        return {"name": job.name, "status": "error", "error": f"failed to write {out_img}"}
    os.replace(tmp, out_img)

    h, w = frame.shape[:2]
    return {"name": job.name, "status": "ok", "frame_index": frame_index,
            "total_frames": total, "width": w, "height": h}


def run_pool(pool, fn, jobs, label, *args):
    results, errors, t0 = {}, [], time.time()
    futures = [pool.submit(fn, j, *args) for j in jobs]
    for n, fut in enumerate(as_completed(futures), 1):
        r = fut.result()
        results[r["name"]] = r
        if r["status"] == "error":
            errors.append(r)
        if n % 250 == 0 or n == len(jobs):
            rate = n / max(time.time() - t0, 1e-6)
            print(f"  {label} {n}/{len(jobs)}  {rate:5.1f}/s  "
                  f"eta {(len(jobs)-n)/max(rate,1e-6)/60:4.1f}m  errors {len(errors)}", flush=True)
    return results, errors


# ── Step 6: write manifest + dataset.json ───────────────────────────────────────
def write_outputs(dst: str, jobs: list[Job], refs: dict[str, dict]) -> None:
    rows = []
    for j in jobs:
        r = refs.get(j.name, {})
        rows.append({
            'style':                    j.style,
            'index':                    j.index,
            'renamed_video':            j.video,
            'reference_image':          j.reference_image,
            'reference_source_video':   j.donor_video,
            'reference_frame_index':    r.get('frame_index', ''),
            'reference_source_frames':  r.get('total_frames', ''),
            'original_video':           j.original_video,
            'original_org':             j.original_org,
            'prompt':                   j.prompt,
        })
    pd.DataFrame(rows).to_csv(os.path.join(dst, "manifest.csv"), index=False)

    dataset = [{"caption": j.prompt, "reference_path": j.reference_image, "media_path": j.video}
               for j in jobs]
    with open(os.path.join(dst, "dataset.json"), "w") as fh:
        json.dump(dataset, fh, indent=2, ensure_ascii=False)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--csv", default=DEFAULT_CSV)
    ap.add_argument("--video-base", default=DEFAULT_VIDEO_BASE)
    ap.add_argument("--dst", default=None,
                    help=f"output dir (default: {DEFAULT_DST}, or "
                         f"{DEFAULT_DST_HORIZONTAL} with --horizontal-only)")
    ap.add_argument("--horizontal-only", action="store_true",
                    help="keep only landscape clips (width > height); references are cut "
                         "from those same clips, so they are landscape too")
    ap.add_argument("--per-style", type=int, default=100, help="samples per style")
    ap.add_argument("--sample-seed", type=int, default=42, help="seed for clip selection")
    ap.add_argument("--seed", type=int, default=20260802, help="seed for donor/frame choice")
    ap.add_argument("--ext", default="png", choices=["png", "jpg"])
    ap.add_argument("--jpeg-quality", type=int, default=95)
    ap.add_argument("--workers", type=int, default=min(32, os.cpu_count() or 8))
    ap.add_argument("--no-clean", action="store_true",
                    help="keep existing files in --dst instead of clearing it first")
    args = ap.parse_args()
    if args.dst is None:
        args.dst = DEFAULT_DST_HORIZONTAL if args.horizontal_only else DEFAULT_DST

    os.makedirs(args.dst, exist_ok=True)
    final_df = select_samples(args.csv, args.video_base, args.per_style, args.sample_seed,
                              horizontal_only=args.horizontal_only, workers=args.workers)
    for style, grp in final_df.groupby('style'):
        print(f"  {style:<30} {len(grp):>4}")

    jobs = plan_jobs(final_df, args.video_base, args.ext, args.seed)
    print(f"\nPlanned {len(jobs)} reference-conditioned samples -> {args.dst}")

    if not args.no_clean:
        stale = [f for f in os.listdir(args.dst) if os.path.isfile(os.path.join(args.dst, f))]
        print(f"Clearing {len(stale)} existing files from {args.dst} ...")
        for f in stale:
            os.remove(os.path.join(args.dst, f))

    t0 = time.time()
    print(f"\nCopying stylized clips ({args.workers} threads)...")
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        copied, copy_errors = run_pool(pool, copy_video, jobs, "copy", args.dst)

    jobs = [j for j in jobs if copied.get(j.name, {}).get("status") != "error"]
    ok_names = {j.name for j in jobs}
    # A dropped clip may still be somebody's donor; those references cannot be cut.
    jobs = [j for j in jobs if j.donor_name in ok_names and j.donor_video]

    print(f"\nExtracting reference frames ({args.workers} processes)...")
    with ProcessPoolExecutor(max_workers=args.workers) as pool:
        refs, ref_errors = run_pool(pool, extract_reference, jobs, "ref",
                                    args.dst, args.seed, args.jpeg_quality,
                                    args.horizontal_only)

    jobs = [j for j in jobs if refs.get(j.name, {}).get("status") != "error"]
    write_outputs(args.dst, jobs, refs)

    errors = copy_errors + ref_errors
    print(f"\nDone in {(time.time()-t0)/60:.1f}m — {len(jobs)} samples written, "
          f"{len(errors)} failed")
    print(f"manifest.csv / dataset.json saved to {args.dst}")
    if errors:
        with open(os.path.join(args.dst, "errors.json"), "w") as fh:
            json.dump(errors, fh, indent=2)
        for e in errors[:10]:
            print(f"  [error] {e['name']}: {e['error']}")
    return 1 if errors else 0


if __name__ == "__main__":
    sys.exit(main())