optimize_assets.py

#!/usr/bin/env python3
"""
optimize_assets.py
------------------
Batch-optimizes all OBJ models and textures under an assets/ directory.

  Models  → decimated via BlenderPythonDecimator (blenderSimplify.py)
            Skipped if already at or below --target-size.
            Ratio is calculated dynamically per model so the result
            lands at exactly --target-size (not a fixed ratio for every file).
  Textures → downscaled via ImageMagick `mogrify`

Usage
-----
python3 optimize_assets.py \
    --target-size 10 \
    --min-ratio  0.1 \
    --dry-run

Arguments
---------
  --target-size   OBJs at or below this size (MB) are left alone.
                  OBJs above it are decimated down to this size. Default: 10
  --min-ratio     Hard floor on the calculated ratio so the mesh is never
                  decimated below this fraction of original faces. Default: 0.1

Requirements
------------
  • Blender installed and on PATH  (or set BLENDER env var)
  • ImageMagick installed           (provides `mogrify`)
  • Python 3.7+

"""

import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path

# ── helpers ───────────────────────────────────────────────────────────────────

os.environ["BLENDER"] = "/Applications/Blender.app/Contents/MacOS/Blender"

def find_blender() -> str:
    """Return the blender executable, preferring the BLENDER env var."""
    return os.environ.get("BLENDER", "blender")


def run(cmd: list[str], dry_run: bool) -> bool:
    """Print and optionally execute a command. Returns True on success."""
    print("  $", " ".join(str(c) for c in cmd))
    if dry_run:
        return True
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"  ✗ FAILED (exit {result.returncode})")
        print(result.stderr[-800:] if result.stderr else "")
        return False
    return True


# ── decimation ────────────────────────────────────────────────────────────────

MB = 1024 * 1024  # bytes per megabyte


def calculate_ratio(current_bytes: int, target_mb: float, min_ratio: float) -> float:
    """
    Derive the decimation ratio needed to bring an OBJ down to target_mb.

    OBJ files are plain text, so byte size scales linearly with face count.
    ratio = target / current  gives us the fraction of faces to keep.
    We clamp to [min_ratio, 1.0] so we never go below the safety floor.
    """
    ratio = (target_mb * MB) / current_bytes
    return max(min_ratio, min(ratio, 1.0))


def decimate_obj(obj_path: Path, decimator_script: Path,
                 target_mb: float, min_ratio: float,
                 blender: str, dry_run: bool) -> None:

    size_bytes = obj_path.stat().st_size
    size_mb    = size_bytes / MB
    label      = f"{obj_path.relative_to(obj_path.parents[2])}"

    if size_mb <= target_mb:
        print(f"\n  [OBJ] {label}")
        print(f"  ✓ skipped  ({size_mb:.1f} MB ≤ {target_mb} MB threshold)")
        return

    ratio = calculate_ratio(size_bytes, target_mb, min_ratio)
    print(f"\n  [OBJ] {label}")
    print(f"  ↓ {size_mb:.1f} MB → target {target_mb} MB  (ratio={ratio:.3f})")

    cmd = [
        blender, "-b", "-P", str(decimator_script),
        "--",
        "--ratio", f"{ratio:.4f}",
        "--inm",   str(obj_path),
        "--outm",  str(obj_path),   # overwrite in-place
    ]
    ok = run(cmd, dry_run)
    if ok and not dry_run:
        new_mb = obj_path.stat().st_size / MB
        print(f"  ✓ decimated  {size_mb:.1f} MB → {new_mb:.1f} MB")


# ── texture downscale ─────────────────────────────────────────────────────────

TEXTURE_EXTS = {".png", ".jpg", ".jpeg", ".tga", ".bmp", ".tiff", ".webp"}

def downscale_texture(tex: Path, target_mb: float, dry_run: bool) -> None:
    import math
    size_bytes = tex.stat().st_size
    size_mb    = size_bytes / MB
    label      = f"{tex.relative_to(tex.parents[2])}"

    print(f"\n  [TEX] {label}")

    if size_mb <= target_mb:
        print(f"  ✓ skipped  ({size_mb:.1f} MB ≤ {target_mb} MB threshold)")
        return

    # File size scales with pixel area (w × h), so to hit the target size we
    # need to scale linear dimensions by sqrt(target / current).
    scale_pct = math.sqrt((target_mb * MB) / size_bytes) * 100
    print(f"  ↓ {size_mb:.1f} MB → target {target_mb} MB  (scale={scale_pct:.1f}%)")

    cmd = [
        "mogrify",
        "-resize", f"{scale_pct:.2f}%",
        str(tex),
    ]
    ok = run(cmd, dry_run)
    if ok and not dry_run:
        new_mb = tex.stat().st_size / MB
        print(f"  ✓ resized  {size_mb:.1f} MB → {new_mb:.1f} MB")


# ── main walk ─────────────────────────────────────────────────────────────────

def process_assets(assets_dir: Path, decimator_script: Path,
                   target_mb: float, min_ratio: float,
                   blender: str, dry_run: bool) -> None:

    model_folders = sorted(
        p for p in assets_dir.iterdir()
        if p.is_dir() and not p.name.startswith("_")
    )

    if not model_folders:
        print(f"No model folders found under {assets_dir}")
        sys.exit(1)

    print(f"\nFound {len(model_folders)} model folder(s) in {assets_dir}\n")
    print("=" * 60)

    for folder in model_folders:
        print(f"\n{'─'*60}")
        print(f"  Folder: {folder.name}")

        # ── OBJ files ─────────────────────────────────────────────────────────
        obj_files = list(folder.glob("*.obj"))
        if not obj_files:
            print("  (no .obj files, skipping decimation)")
        for obj in obj_files:
            decimate_obj(obj, decimator_script, target_mb, min_ratio, blender, dry_run)

        # ── Texture files ──────────────────────────────────────────────────────
        tex_files = [
            f for f in folder.iterdir()
            if f.is_file() and f.suffix.lower() in TEXTURE_EXTS
        ]
        if not tex_files:
            print("  (no texture files, skipping downscale)")
        for tex in tex_files:
            downscale_texture(tex, target_mb, dry_run)

    print(f"\n{'='*60}")
    print("Done." if not dry_run else "Dry-run complete — no files were changed.")


# ── CLI ───────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Batch-optimize OBJ models and textures.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--target-size", type=float, default=10.0,
        help="Target OBJ size in MB. Files at or below this are skipped. Default: 10",
    )
    parser.add_argument(
        "--min-ratio", type=float, default=0.1,
        help="Hard floor on decimation ratio — never reduce below this fraction "
             "of original faces, even if the file is very large. Default: 0.1",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Print commands without executing them.",
    )
    args = parser.parse_args()

    assets_dir = Path("assets")
    decimator_script = Path("blenderSimplify.py")
    blender = find_blender()

    # ── pre-flight checks ─────────────────────────────────────────────────────
    errors = []

    if not assets_dir.is_dir():
        errors.append(f"Assets directory not found: {assets_dir}")

    if not decimator_script.is_file():
        errors.append(f"Decimator script not found: {decimator_script}")

    if args.target_size <= 0:
        errors.append("--target-size must be > 0")

    if not (0.05 <= args.min_ratio <= 1.0):
        errors.append("--min-ratio must be between 0.05 and 1.0")

    if not shutil.which(blender):
        errors.append(
            f"Blender not found on PATH (tried '{blender}'). "
            "Install Blender or set the BLENDER environment variable."
        )

    if not shutil.which("mogrify"):
        errors.append(
            "ImageMagick `mogrify` not found on PATH. "
            "Install ImageMagick:  brew install imagemagick  "
            "or  sudo apt install imagemagick"
        )

    if errors:
        print("\nPre-flight errors:")
        for e in errors:
            print(f"  ✗ {e}")
        sys.exit(1)

    # ── summary ───────────────────────────────────────────────────────────────
    print("\noptimize_assets.py")
    print(f"  assets dir   : {assets_dir}")
    print(f"  decimator    : {decimator_script}")
    print(f"  blender      : {blender}")
    print(f"  target size  : {args.target_size} MB  (skip if already ≤ this)")
    print(f"  min ratio    : {args.min_ratio}  (hard floor on face reduction)")
    print(f"  dry run      : {args.dry_run}")

    process_assets(
        assets_dir, decimator_script,
        args.target_size, args.min_ratio,
        blender, args.dry_run,
    )


if __name__ == "__main__":
    main()