generate_validation.py

import os
import re
import shutil
import pandas as pd

VIDEO_BASE  = "/mnt/data/xinyuy/datasets/OpenVE-3M/videos"
CSV_PATH    = "/mnt/data/xinyuy/datasets/OpenVE-3M/csv_files/global_style.csv"
OUT_DIR     = "/home/xinyuy/dataset_processing/OpenVE-3M/Style-Tranfer-Validation"

EXCLUDE_MANIFESTS = [
    "/home/xinyuy/dataset_processing/OpenVE-3M/Style-Transfer-OpenVE-3M-Dataset/manifest.csv",
    "/home/xinyuy/dataset_processing/OpenVE-3M/Style-Tranfer/manifest.csv",
]

N_SAMPLES = 10

# ── Step 1: collect all video hashes already used in existing datasets ─────────
print("Loading exclusion manifests...")
used_videos = set()
for path in EXCLUDE_MANIFESTS:
    m = pd.read_csv(path)
    used_videos.update(m['original_video'].tolist())
    used_videos.update(m['original_org'].tolist())
print(f"Excluding {len(used_videos):,} already-used video paths")

# ── Step 2: load CSV and keep rows where both videos exist ─────────────────────
print("Loading CSV...")
data = pd.read_csv(CSV_PATH, sep=';')

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

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

data['video_new']          = data['video'].apply(remap)
data['original_video_new'] = data['original_video'].apply(remap)

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

# ── Step 3: extract style, normalize, filter Apply-only + CANONICAL_19 ────────
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|filter|effect|technique|art style|principles|to this)|\s*,|\s*\.)'

def extract_style(prompt):
    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

_NORMALIZE = {
    'dynamic':                     'Abstract Art',
    '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',
}

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

# Forward "Apply" rows, canonical styles, not already used
gs_common = gs_full[
    gs_full['style'].isin(CANONICAL_19) &
    gs_full['prompt'].str.startswith('Apply') &
    ~gs_full['video'].isin(used_videos) &
    ~gs_full['original_video'].isin(used_videos)
]
print(f"Available (unseen) forward rows: {len(gs_common):,}")

# ── Step 4: randomly sample N_SAMPLES rows ────────────────────────────────────
sampled = gs_common.sample(n=N_SAMPLES, random_state=7).reset_index(drop=True)

print(f"\nSampled {len(sampled)} rows:")
for _, row in sampled.iterrows():
    print(f"  [{row['style']:<28}] {row['prompt'][:60]}...")

# ── Step 5: create output dir and copy files ──────────────────────────────────
os.makedirs(OUT_DIR, exist_ok=True)

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',
}

def resolve(rel_path):
    remapped = rel_path.replace('global_style/', 'global_style_new/', 1)
    full = os.path.join(VIDEO_BASE, remapped)
    if os.path.exists(full):
        return full
    fallback = os.path.join(VIDEO_BASE, rel_path)
    return fallback if os.path.exists(fallback) else None

missing, copied = 0, 0
mapping_rows = []

print("\nCopying files...")
for idx, row in sampled.iterrows():
    prefix = STYLE_PREFIX[row['style']]
    stem   = f"{prefix}_{idx:03d}"

    src_edited   = resolve(row['video'])
    src_original = resolve(row['original_video'])
    dst_edited   = os.path.join(OUT_DIR, f"{stem}.mp4")
    dst_original = os.path.join(OUT_DIR, f"{stem}_org.mp4")

    for src, dst in [(src_edited, dst_edited), (src_original, dst_original)]:
        if src is None:
            print(f"  MISSING: {row['video'] if dst == dst_edited else row['original_video']}")
            missing += 1
        else:
            shutil.copy2(src, dst)
            copied += 1

    mapping_rows.append({
        'style':          row['style'],
        'index':          idx,
        'renamed_video':  os.path.basename(dst_edited),
        'renamed_org':    os.path.basename(dst_original),
        'original_video': row['video'],
        'original_org':   row['original_video'],
        'prompt':         row['prompt'],
    })

mapping_df = pd.DataFrame(mapping_rows)
mapping_df.to_csv(os.path.join(OUT_DIR, "manifest.csv"), index=False)

print(f"\nDone. Copied: {copied}  |  Missing: {missing}")
print(f"Manifest saved to {OUT_DIR}/manifest.csv")
print("\nFiles in validation folder:")
for f in sorted(os.listdir(OUT_DIR)):
    print(f"  {f}")