438 lines
17 KiB
Python
438 lines
17 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Compress PNGs in unchanged render sequences. Writes to staging, prompts for Y/N, then overwrites originals."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import multiprocessing
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
import shutil
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
try:
|
||
|
|
from PIL import Image
|
||
|
|
except ImportError:
|
||
|
|
print("Error: Pillow (PIL) is required. Install with: pip install Pillow", file=sys.stderr)
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
try:
|
||
|
|
import unlock_processpool
|
||
|
|
unlock_processpool.please()
|
||
|
|
UNLOCKED = True
|
||
|
|
except ImportError:
|
||
|
|
UNLOCKED = False
|
||
|
|
|
||
|
|
RENDER_ROOT = Path("Renders")
|
||
|
|
ARCHIVE_ROOT = RENDER_ROOT / "_zipped"
|
||
|
|
STAGING_DIR = "_compressed_staging"
|
||
|
|
CORRUPTED_DIR = "corrupted"
|
||
|
|
STATE_SUFFIX = ".meta.json"
|
||
|
|
PNG_EXT = ".png"
|
||
|
|
|
||
|
|
|
||
|
|
def is_archive_path(path: Path) -> bool:
|
||
|
|
return any(part in ("_archive", "_CURRENT", "_zipped", STAGING_DIR, CORRUPTED_DIR) for part in path.parts)
|
||
|
|
|
||
|
|
|
||
|
|
def find_sequence_dirs(root: Path) -> list[Path]:
|
||
|
|
seen = set()
|
||
|
|
out = []
|
||
|
|
for dirpath, dirnames, filenames in os.walk(root):
|
||
|
|
path = Path(dirpath)
|
||
|
|
dirnames[:] = [d for d in dirnames if d not in ("_archive", "_CURRENT", STAGING_DIR, CORRUPTED_DIR, "_zipped")]
|
||
|
|
if is_archive_path(path):
|
||
|
|
continue
|
||
|
|
if any(Path(dirpath, f).suffix.lower() == PNG_EXT for f in filenames):
|
||
|
|
resolved = path.resolve()
|
||
|
|
if resolved not in seen:
|
||
|
|
seen.add(resolved)
|
||
|
|
out.append(path)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def iter_png_files(seq_dir: Path):
|
||
|
|
for dirpath, dirnames, filenames in os.walk(seq_dir):
|
||
|
|
path = Path(dirpath)
|
||
|
|
dirnames[:] = [d for d in dirnames if d not in ("_archive", "_CURRENT", "_zipped", STAGING_DIR, CORRUPTED_DIR)]
|
||
|
|
if is_archive_path(path):
|
||
|
|
continue
|
||
|
|
for f in filenames:
|
||
|
|
if Path(f).suffix.lower() == PNG_EXT:
|
||
|
|
yield path / f
|
||
|
|
|
||
|
|
|
||
|
|
def load_state(state_path: Path) -> dict | None:
|
||
|
|
if not state_path.exists():
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
state = json.loads(state_path.read_text())
|
||
|
|
if "files" in state:
|
||
|
|
state["files"] = [e for e in state.get("files", []) if Path(e.get("path", "")).name.lower() != "thumbs.db"]
|
||
|
|
if platform.system() == "Windows" and "files" in state:
|
||
|
|
for e in state.get("files", []):
|
||
|
|
if "mtime_ns" in e:
|
||
|
|
e["mtime_ns"] = (e["mtime_ns"] // 100) * 100
|
||
|
|
return state
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def state_changed(seq_state: dict, stored_state: dict | None) -> bool:
|
||
|
|
if stored_state is None:
|
||
|
|
return True
|
||
|
|
is_windows = platform.system() == "Windows"
|
||
|
|
|
||
|
|
def png_only(files):
|
||
|
|
return [e for e in files if Path(e.get("path", "")).suffix.lower() == ".png"]
|
||
|
|
|
||
|
|
def norm(s):
|
||
|
|
out = []
|
||
|
|
for e in png_only(s.get("files", [])):
|
||
|
|
if Path(e.get("path", "")).name.lower() == "thumbs.db":
|
||
|
|
continue
|
||
|
|
m = e.get("mtime_ns", 0)
|
||
|
|
if is_windows:
|
||
|
|
m = (m // 100) * 100
|
||
|
|
out.append({"path": e["path"], "size": e["size"], "mtime_ns": m})
|
||
|
|
return {"files": out}
|
||
|
|
|
||
|
|
return norm(seq_state) != norm(stored_state)
|
||
|
|
|
||
|
|
|
||
|
|
def compute_state(seq_dir: Path) -> dict:
|
||
|
|
entries = []
|
||
|
|
is_windows = platform.system() == "Windows"
|
||
|
|
for p in sorted(iter_png_files(seq_dir), key=lambda x: x.relative_to(seq_dir).as_posix()):
|
||
|
|
stat = p.stat()
|
||
|
|
mtime_ns = stat.st_mtime_ns
|
||
|
|
if is_windows:
|
||
|
|
mtime_ns = (mtime_ns // 100) * 100
|
||
|
|
entries.append({"path": p.relative_to(seq_dir).as_posix(), "size": stat.st_size, "mtime_ns": mtime_ns})
|
||
|
|
return {"files": entries}
|
||
|
|
|
||
|
|
|
||
|
|
def find_archive_and_state(seq_dir: Path) -> tuple[Path | None, Path | None]:
|
||
|
|
"""Return (archive_path, state_path) if sequence is archived and unchanged; else (None, None)."""
|
||
|
|
try:
|
||
|
|
rel = seq_dir.relative_to(RENDER_ROOT)
|
||
|
|
except ValueError:
|
||
|
|
return None, None
|
||
|
|
for suffix in (".7z", ".zip"):
|
||
|
|
archive = ARCHIVE_ROOT / f"{rel}{suffix}"
|
||
|
|
if archive.exists():
|
||
|
|
state = archive.with_suffix(archive.suffix + STATE_SUFFIX)
|
||
|
|
return archive, state
|
||
|
|
return None, None
|
||
|
|
|
||
|
|
|
||
|
|
def compress_png(input_path: Path, output_path: Path, force_bitdepth: str | None) -> tuple[str, bool, str | None, int, int, float, bool, bool]:
|
||
|
|
"""Returns (path, success, error, orig_size, new_size, savings_pct, was_skipped, is_corrupted)."""
|
||
|
|
try:
|
||
|
|
output_path = Path(output_path)
|
||
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
orig_size = input_path.stat().st_size
|
||
|
|
if orig_size == 0:
|
||
|
|
return (str(input_path), False, "CORRUPTED: File is 0 bytes (empty/placeholder)", 0, 0, 0, False, True)
|
||
|
|
try:
|
||
|
|
img = Image.open(input_path)
|
||
|
|
img.load()
|
||
|
|
except Exception as e:
|
||
|
|
return (str(input_path), False, f"CORRUPTED: Cannot open/load image - {e}", orig_size, 0, 0, False, True)
|
||
|
|
if img.width == 0 or img.height == 0:
|
||
|
|
return (str(input_path), False, f"CORRUPTED: Invalid dimensions ({img.width}x{img.height})", orig_size, 0, 0, False, True)
|
||
|
|
|
||
|
|
# All-black + suspiciously small (corruption indicator)
|
||
|
|
try:
|
||
|
|
sample_size = min(100, img.width * img.height)
|
||
|
|
pixels = list(img.getdata()[:sample_size])
|
||
|
|
if pixels:
|
||
|
|
all_black = True
|
||
|
|
for pixel in pixels:
|
||
|
|
if isinstance(pixel, (tuple, list)):
|
||
|
|
if any(p > 0 for p in pixel):
|
||
|
|
all_black = False
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
if pixel > 0:
|
||
|
|
all_black = False
|
||
|
|
break
|
||
|
|
if all_black and orig_size < (img.width * img.height * 0.1):
|
||
|
|
return (str(input_path), False, "CORRUPTED: Image appears all black with suspiciously small file size", orig_size, 0, 0, False, True)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
# Determine target bit depth
|
||
|
|
if force_bitdepth == "8":
|
||
|
|
is_16bit = False
|
||
|
|
elif force_bitdepth == "16":
|
||
|
|
is_16bit = True
|
||
|
|
else:
|
||
|
|
is_16bit = img.mode == "I"
|
||
|
|
if not is_16bit and img.mode in ("RGB", "RGBA", "LA"):
|
||
|
|
try:
|
||
|
|
pixels = list(img.getdata())
|
||
|
|
if pixels:
|
||
|
|
if isinstance(pixels[0], (tuple, list)):
|
||
|
|
max_val = max(max(p) for p in pixels[:1000])
|
||
|
|
else:
|
||
|
|
max_val = max(pixels[:1000])
|
||
|
|
if max_val > 255:
|
||
|
|
is_16bit = True
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
# Handle bit depth conversion (8-bit to 8-bit, 16-bit to 16-bit)
|
||
|
|
if is_16bit:
|
||
|
|
if force_bitdepth == "16":
|
||
|
|
if img.mode != "I":
|
||
|
|
img = img.convert("RGBA" if ("A" in img.mode or img.mode == "LA") else "RGB")
|
||
|
|
else:
|
||
|
|
if img.mode == "I":
|
||
|
|
pass
|
||
|
|
elif img.mode in ("RGB", "RGBA"):
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
img = img.convert("RGBA" if "A" in img.mode else "RGB")
|
||
|
|
else:
|
||
|
|
if force_bitdepth == "8":
|
||
|
|
if img.mode == "I":
|
||
|
|
img = img.convert("L")
|
||
|
|
elif img.mode == "RGBA":
|
||
|
|
pass
|
||
|
|
elif img.mode == "RGB":
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
img = img.convert("RGBA" if "A" in img.mode else "RGB")
|
||
|
|
else:
|
||
|
|
if img.mode == "RGBA":
|
||
|
|
pass
|
||
|
|
elif img.mode != "RGB":
|
||
|
|
img = img.convert("RGB")
|
||
|
|
|
||
|
|
img.save(str(output_path), "PNG", optimize=True, compress_level=9)
|
||
|
|
new_size = output_path.stat().st_size
|
||
|
|
savings = (orig_size - new_size) / orig_size * 100 if orig_size > 0 else 0
|
||
|
|
return (str(input_path), True, None, orig_size, new_size, savings, False, False)
|
||
|
|
except Exception as e:
|
||
|
|
corrupt = "truncated" in str(e).lower() or "cannot identify" in str(e).lower() or "corrupt" in str(e).lower()
|
||
|
|
return (str(input_path), False, str(e), 0, 0, 0, False, corrupt)
|
||
|
|
|
||
|
|
|
||
|
|
def move_to_corrupted(src: Path, base: Path, corrupted_root: Path) -> bool:
|
||
|
|
try:
|
||
|
|
rel = src.relative_to(base)
|
||
|
|
except ValueError:
|
||
|
|
rel = src.name
|
||
|
|
dest = corrupted_root / rel
|
||
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
src.rename(dest)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def format_size(size_bytes: float) -> str:
|
||
|
|
for unit in ("B", "KB", "MB", "GB"):
|
||
|
|
if size_bytes < 1024.0:
|
||
|
|
return f"{size_bytes:.2f} {unit}"
|
||
|
|
size_bytes /= 1024.0
|
||
|
|
return f"{size_bytes:.2f} TB"
|
||
|
|
|
||
|
|
|
||
|
|
def format_time(seconds: float) -> str:
|
||
|
|
hours = int(seconds // 3600)
|
||
|
|
minutes = int((seconds % 3600) // 60)
|
||
|
|
secs = int(seconds % 60)
|
||
|
|
centiseconds = int((seconds % 1) * 100)
|
||
|
|
return f"{hours:02d}:{minutes:02d}:{secs:02d}:{centiseconds:02d}"
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser(description="Compress PNGs in unchanged sequences; staging + Y/N overwrite.")
|
||
|
|
parser.add_argument("--8bit", "-8", action="store_true", dest="force_8bit")
|
||
|
|
parser.add_argument("--16bit", "-16", action="store_true", dest="force_16bit")
|
||
|
|
parser.add_argument("--verbose", action="store_true")
|
||
|
|
args = parser.parse_args()
|
||
|
|
force_bitdepth = "8" if args.force_8bit else ("16" if args.force_16bit else None)
|
||
|
|
if args.force_8bit and args.force_16bit:
|
||
|
|
print("Error: Cannot specify both --8bit and --16bit. Choose one.")
|
||
|
|
return 1
|
||
|
|
if force_bitdepth == "8":
|
||
|
|
print("Mode: Forcing 8-bit color depth")
|
||
|
|
elif force_bitdepth == "16":
|
||
|
|
print("Mode: Forcing 16-bit color depth")
|
||
|
|
else:
|
||
|
|
print("Mode: Auto-detect bit depth (preserve 16-bit if present)")
|
||
|
|
|
||
|
|
if not RENDER_ROOT.exists():
|
||
|
|
print("[ERROR] Renders directory not found.")
|
||
|
|
return 1
|
||
|
|
|
||
|
|
staging_root = RENDER_ROOT / STAGING_DIR
|
||
|
|
corrupted_root = RENDER_ROOT / CORRUPTED_DIR
|
||
|
|
|
||
|
|
# Find changed sequences (no archive, or state differs from stored)
|
||
|
|
changed: list[Path] = []
|
||
|
|
for seq_dir in find_sequence_dirs(RENDER_ROOT):
|
||
|
|
current = compute_state(seq_dir)
|
||
|
|
if not current["files"]:
|
||
|
|
continue
|
||
|
|
_, state_path = find_archive_and_state(seq_dir)
|
||
|
|
if state_path is None or not state_path.exists():
|
||
|
|
changed.append(seq_dir) # New sequence, never zipped
|
||
|
|
continue
|
||
|
|
stored = load_state(state_path)
|
||
|
|
if state_changed(current, stored):
|
||
|
|
changed.append(seq_dir)
|
||
|
|
|
||
|
|
if not changed:
|
||
|
|
print("No changed sequences with PNGs found. All sequences are up to date.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
print(f"Found {len(changed)} changed sequence(s) to compress.")
|
||
|
|
staging_root.mkdir(parents=True, exist_ok=True)
|
||
|
|
corrupted_root.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
# Collect PNGs to compress
|
||
|
|
work: list[tuple[Path, Path]] = []
|
||
|
|
for seq_dir in changed:
|
||
|
|
rel = seq_dir.relative_to(RENDER_ROOT)
|
||
|
|
staging_seq = staging_root / rel
|
||
|
|
for png in iter_png_files(seq_dir):
|
||
|
|
out = staging_seq / png.relative_to(seq_dir)
|
||
|
|
work.append((png, out))
|
||
|
|
|
||
|
|
if not work:
|
||
|
|
print("No PNG files in changed sequences.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
max_workers = multiprocessing.cpu_count()
|
||
|
|
if platform.system() == "Windows" and not UNLOCKED:
|
||
|
|
max_workers = min(max_workers, 61)
|
||
|
|
if max_workers < multiprocessing.cpu_count():
|
||
|
|
print(f"Detected {multiprocessing.cpu_count()} CPU threads, but Windows limits ProcessPoolExecutor to 61 workers.")
|
||
|
|
print("Install 'unlock-processpool-win' package to use all cores: pip install unlock-processpool-win")
|
||
|
|
else:
|
||
|
|
if UNLOCKED:
|
||
|
|
print("Using unlock-processpool-win to bypass Windows 61-worker limit")
|
||
|
|
print(f"Compressing {len(work)} PNGs with {max_workers} workers...")
|
||
|
|
print("-" * 80)
|
||
|
|
|
||
|
|
compressed = 0
|
||
|
|
corrupted_count = 0
|
||
|
|
corrupted_paths: list[tuple[str, str]] = []
|
||
|
|
total_orig = 0
|
||
|
|
total_new = 0
|
||
|
|
failed = 0
|
||
|
|
start_time = time.time()
|
||
|
|
last_update_time = start_time
|
||
|
|
|
||
|
|
with ProcessPoolExecutor(max_workers=max_workers) as ex:
|
||
|
|
futures = {ex.submit(compress_png, inp, out, force_bitdepth): (inp, out) for inp, out in work}
|
||
|
|
for future in as_completed(futures):
|
||
|
|
inp, out = futures[future]
|
||
|
|
path, ok, err, o_sz, n_sz, _, _, is_corrupt = future.result()
|
||
|
|
if ok:
|
||
|
|
compressed += 1
|
||
|
|
total_orig += o_sz
|
||
|
|
total_new += n_sz
|
||
|
|
elif is_corrupt:
|
||
|
|
corrupted_count += 1
|
||
|
|
try:
|
||
|
|
if move_to_corrupted(Path(path), RENDER_ROOT, corrupted_root):
|
||
|
|
corrupted_paths.append((path, f"Moved to {corrupted_root / Path(path).relative_to(RENDER_ROOT)}"))
|
||
|
|
else:
|
||
|
|
corrupted_paths.append((path, str(err)))
|
||
|
|
except Exception as e:
|
||
|
|
corrupted_paths.append((path, str(e)))
|
||
|
|
print(f"\n[CORRUPTED] {path} -> {err}")
|
||
|
|
else:
|
||
|
|
failed += 1
|
||
|
|
print(f"\n[ERROR] {path}: {err}")
|
||
|
|
|
||
|
|
# Live progress (update every 0.5s or on first completion)
|
||
|
|
current_time = time.time()
|
||
|
|
time_since_update = current_time - last_update_time
|
||
|
|
processed = compressed + corrupted_count + failed
|
||
|
|
if processed == 1 or time_since_update >= 0.5:
|
||
|
|
elapsed = current_time - start_time
|
||
|
|
rate = processed / elapsed if elapsed > 0 else 0
|
||
|
|
remaining = len(work) - processed
|
||
|
|
eta_seconds = remaining / rate if rate > 0 and remaining > 0 else 0
|
||
|
|
total_savings = total_orig - total_new
|
||
|
|
total_savings_pct = (total_savings / total_orig * 100) if total_orig > 0 else 0
|
||
|
|
eta_str = format_time(eta_seconds) if eta_seconds > 0 else "calculating..."
|
||
|
|
elapsed_str = format_time(elapsed)
|
||
|
|
print(
|
||
|
|
f"[{processed:5d}/{len(work)}] "
|
||
|
|
f"Compressed: {compressed} | Corrupted: {corrupted_count} | Failed: {failed} | "
|
||
|
|
f"Speed: {rate:.1f} files/sec | "
|
||
|
|
f"Elapsed: {elapsed_str} | ETA: {eta_str} | "
|
||
|
|
f"Saved: {format_size(total_savings)} ({total_savings_pct:.1f}%)",
|
||
|
|
end="\r",
|
||
|
|
)
|
||
|
|
last_update_time = current_time
|
||
|
|
|
||
|
|
print() # Clear progress line
|
||
|
|
total_time = time.time() - start_time
|
||
|
|
total_savings = total_orig - total_new
|
||
|
|
total_savings_pct = (total_savings / total_orig * 100) if total_orig > 0 else 0
|
||
|
|
avg_rate = (compressed + corrupted_count + failed) / total_time if total_time > 0 else 0
|
||
|
|
|
||
|
|
print("=" * 80)
|
||
|
|
print("Compression complete!")
|
||
|
|
print(f"Successfully compressed: {compressed} files")
|
||
|
|
if corrupted_count > 0:
|
||
|
|
print(f"Corrupted (bad PNGs): {corrupted_count} files")
|
||
|
|
if failed > 0:
|
||
|
|
print(f"Failed: {failed} files")
|
||
|
|
print(f"Total time: {format_time(total_time)}")
|
||
|
|
print(f"Average speed: {avg_rate:.2f} files/second")
|
||
|
|
print(f"Original size: {format_size(total_orig)}")
|
||
|
|
print(f"Compressed size: {format_size(total_new)}")
|
||
|
|
print(f"Total savings: {format_size(total_savings)} ({total_savings_pct:.1f}%)")
|
||
|
|
|
||
|
|
if corrupted_paths:
|
||
|
|
print(f"\nCorrupted files moved to {corrupted_root}")
|
||
|
|
print("\n" + "=" * 80)
|
||
|
|
print("CORRUPTED FILES LIST:")
|
||
|
|
print("=" * 80)
|
||
|
|
for item in corrupted_paths:
|
||
|
|
file_path, move_status = item
|
||
|
|
print(f" {file_path}")
|
||
|
|
print(f" Status: {move_status}")
|
||
|
|
|
||
|
|
check_dirs = [f"Renders\\{STAGING_DIR}"]
|
||
|
|
if corrupted_count > 0:
|
||
|
|
check_dirs.append(f"Renders\\{CORRUPTED_DIR}")
|
||
|
|
resp = input(f"\nCheck files in {' and '.join(check_dirs)}. Ready to overwrite originals? (Y/N): ").strip().upper()
|
||
|
|
if resp != "Y":
|
||
|
|
print("Exiting. Staging and corrupted folders left for inspection.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
# Overwrite originals with staged files
|
||
|
|
for seq_dir in changed:
|
||
|
|
rel = seq_dir.relative_to(RENDER_ROOT)
|
||
|
|
staging_seq = staging_root / rel
|
||
|
|
if not staging_seq.exists():
|
||
|
|
continue
|
||
|
|
for staged in staging_seq.rglob("*.png"):
|
||
|
|
rel_file = staged.relative_to(staging_seq)
|
||
|
|
orig = seq_dir / rel_file
|
||
|
|
if orig.exists():
|
||
|
|
shutil.copy2(staged, orig)
|
||
|
|
|
||
|
|
# Cleanup staging and corrupted
|
||
|
|
shutil.rmtree(staging_root, ignore_errors=True)
|
||
|
|
if corrupted_root.exists():
|
||
|
|
shutil.rmtree(corrupted_root, ignore_errors=True)
|
||
|
|
print("Overwrite complete. Staging and corrupted folders cleaned up.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|