openroller/tools/dds_to_png.py

#!/usr/bin/env python3
"""
Convert all .dds files under a directory to .png into a single output folder.

We intentionally use ffmpeg for decoding DDS, since ImageMagick often lacks a DDS delegate.

Example:
  python tools/dds_to_png.py GC/data /tmp/gc_dds_png --jobs 8
"""

from __future__ import annotations

import argparse
import hashlib
import os
import shutil
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Tuple


@dataclass(frozen=True)
class Task:
    src: Path
    dst: Path


def _iter_dds(root: Path) -> Iterable[Path]:
    for dirpath, _, filenames in os.walk(root):
        for fn in filenames:
            if fn.lower().endswith(".dds"):
                yield Path(dirpath) / fn


def _safe_flat_name(root: Path, p: Path) -> str:
    rel = p.relative_to(root).as_posix()
    # Flatten path to a filename. Keep ASCII-ish and avoid huge names.
    flat = rel.replace("/", "__").replace("\\", "__")
    if len(flat) > 180:
        h = hashlib.sha1(rel.encode("utf-8")).hexdigest()[:10]
        base = Path(flat).stem[:120]
        flat = f"{base}__{h}.dds"
    return Path(flat).with_suffix(".png").name


def _run_ffmpeg(src: Path, dst: Path) -> Tuple[bool, str]:
    # -frames:v 1 ensures we only keep the first image if ffmpeg treats it as a sequence.
    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(src),
        "-frames:v",
        "1",
        str(dst),
    ]
    try:
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    except FileNotFoundError:
        return False, "ffmpeg not found"
    if p.returncode != 0:
        msg = (p.stderr or p.stdout or "").strip()
        if not msg:
            msg = f"ffmpeg failed with code {p.returncode}"
        return False, msg
    return True, ""


def _convert_one(t: Task, overwrite: bool) -> Tuple[bool, str]:
    if t.dst.exists():
        if not overwrite:
            return True, "skip"
        try:
            t.dst.unlink()
        except Exception as e:
            return False, f"unlink failed: {e}"
    t.dst.parent.mkdir(parents=True, exist_ok=True)
    ok, err = _run_ffmpeg(t.src, t.dst)
    if not ok:
        return False, err
    return True, ""


def main(argv: List[str]) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("input_dir", help="Root directory to scan for .dds (e.g. GC/data)")
    ap.add_argument("output_dir", help="Output directory for flat .png files")
    ap.add_argument("--jobs", type=int, default=8, help="Parallel jobs (default: 8)")
    ap.add_argument("--overwrite", action="store_true", help="Overwrite existing .png")
    ap.add_argument("--max", type=int, default=0, help="Only convert first N files (0 = all)")
    ap.add_argument("--list", action="store_true", help="List discovered files and exit")
    args = ap.parse_args(argv)

    root = Path(args.input_dir).resolve()
    out = Path(args.output_dir).resolve()

    if not root.exists():
        print(f"input_dir does not exist: {root}", file=sys.stderr)
        return 2
    if shutil.which("ffmpeg") is None:
        print("ffmpeg not found in PATH", file=sys.stderr)
        return 2

    sources = sorted(_iter_dds(root))
    if args.max and args.max > 0:
        sources = sources[: args.max]

    if args.list:
        for p in sources:
            print(p)
        return 0

    tasks: List[Task] = []
    used = set()
    collisions = 0
    for p in sources:
        name = _safe_flat_name(root, p)
        if name in used:
            collisions += 1
            h = hashlib.sha1(str(p).encode("utf-8")).hexdigest()[:10]
            name = Path(name).with_suffix("").name + f"__{h}.png"
        used.add(name)
        tasks.append(Task(src=p, dst=out / name))

    print(f"Found {len(tasks)} DDS files under {root}")
    if collisions:
        print(f"Name collisions: {collisions} (resolved with hashes)")
    print(f"Output dir: {out}")

    ok = 0
    fail = 0
    skipped = 0
    errors: List[Tuple[Path, str]] = []

    jobs = max(1, int(args.jobs))
    with ThreadPoolExecutor(max_workers=jobs) as ex:
        futs = {ex.submit(_convert_one, t, args.overwrite): t for t in tasks}
        done = 0
        for f in as_completed(futs):
            t = futs[f]
            done += 1
            try:
                success, msg = f.result()
            except Exception as e:
                success, msg = False, f"exception: {e}"
            if success:
                if msg == "skip":
                    skipped += 1
                else:
                    ok += 1
            else:
                fail += 1
                errors.append((t.src, msg))
            if done % 200 == 0 or done == len(tasks):
                print(f"Progress: {done}/{len(tasks)} ok={ok} skip={skipped} fail={fail}")

    if errors:
        log = out / "_errors.txt"
        out.mkdir(parents=True, exist_ok=True)
        with log.open("w", encoding="utf-8") as fp:
            for src, msg in errors[:2000]:
                fp.write(f"{src}\t{msg}\n")
        print(f"Failures: {fail} (see {log})")

    print(f"Done: ok={ok} skip={skipped} fail={fail}")
    return 0 if fail == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))