make_display.py

"""
For each row in local_add.csv and local_remove.csv where both the edited video
and the original video exist in the videos folder, extract the first frame of
each and save a side-by-side image (original left, edited right) to display/.
"""

import csv
import cv2
import numpy as np
from pathlib import Path

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

# CSV path prefix → actual folder name under VIDEOS_DIR
PREFIX_MAP = {
    "local_add": "local_add",
    "local_remove": "local_remove",   # may not exist
    "global_style": "global_style_new",
    "background_change": "background_change",  # may not exist
    "local_change": "local_change",   # may not exist
}

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


def resolve_path(csv_path: str) -> Path | None:
    """Map a CSV-relative path like 'local_add/abc.mp4' to an absolute path."""
    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
    full_path = VIDEOS_DIR / actual_folder / filename
    return full_path


def get_first_frame(video_path: Path) -> np.ndarray | None:
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        return None
    ok, frame = cap.read()
    cap.release()
    if not ok:
        return None
    return frame  # BGR


def collect_pairs():
    """Return list of (edited_path, original_path, label) tuples."""
    pairs = []
    seen = set()

    for csv_file in CSV_FILES:
        source = csv_file.stem
        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()
                if not edited_csv or not original_csv:
                    continue

                edited_path = resolve_path(edited_csv)
                original_path = resolve_path(original_csv)

                if edited_path is None or original_path is None:
                    continue
                if not edited_path.exists() or not original_path.exists():
                    continue

                key = (str(edited_path), str(original_path))
                if key in seen:
                    continue
                seen.add(key)

                pairs.append((edited_path, original_path, source))

    return pairs


def make_side_by_side(original_frame: np.ndarray, edited_frame: np.ndarray) -> np.ndarray:
    """Resize both frames to the same height then concatenate horizontally."""
    h = max(original_frame.shape[0], edited_frame.shape[0])

    def resize_to_height(img, target_h):
        oh, ow = img.shape[:2]
        scale = target_h / oh
        new_w = int(ow * scale)
        return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)

    orig_resized = resize_to_height(original_frame, h)
    edit_resized = resize_to_height(edited_frame, h)

    # Add a thin white divider
    divider = np.ones((h, 4, 3), dtype=np.uint8) * 255
    return np.concatenate([orig_resized, divider, edit_resized], axis=1)


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

    print("Scanning CSV files for valid pairs...", flush=True)
    pairs = collect_pairs()
    print(f"Found {len(pairs)} valid pairs where both videos exist.", flush=True)

    if not pairs:
        print("No pairs found. Exiting.")
        return

    ok_count = 0
    skip_count = 0

    for i, (edited_path, original_path, source) in enumerate(pairs):
        out_name = edited_path.stem + "__" + original_path.stem + ".jpg"
        out_path = OUTPUT_DIR / out_name

        if out_path.exists():
            ok_count += 1
            continue

        orig_frame = get_first_frame(original_path)
        if orig_frame is None:
            print(f"  [WARN] Could not read frame from {original_path}", flush=True)
            skip_count += 1
            continue

        edit_frame = get_first_frame(edited_path)
        if edit_frame is None:
            print(f"  [WARN] Could not read frame from {edited_path}", flush=True)
            skip_count += 1
            continue

        combined = make_side_by_side(orig_frame, edit_frame)
        cv2.imwrite(str(out_path), combined, [cv2.IMWRITE_JPEG_QUALITY, 90])
        ok_count += 1

        if (i + 1) % 100 == 0:
            print(f"  Processed {i + 1}/{len(pairs)} pairs...", flush=True)

    print(f"\nDone. Saved {ok_count} images to {OUTPUT_DIR}  (skipped {skip_count}).")


if __name__ == "__main__":
    main()