#!/usr/bin/env python3 """Re-segment and compare existing image candidates entirely offline. Original generated sheets, parts, layouts and previews are never modified. Each candidate gets a ``resegmented/`` evaluation directory, and every batch gets a comparison sheet plus report. No image generation API is called. """ from __future__ import annotations import argparse from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timezone import hashlib import json import math from pathlib import Path import re import shutil import subprocess import sys import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont ROOT = Path(__file__).resolve().parent DEFAULT_BATCHES = ( "cowboy-6", "female-adventurer-3", "male-adventurer-3", "male-person-3", "robot-3", "zombie-3", "female-person-3", ) DEFAULT_GENIELABS = Path("/Users/zhaoziyu/Documents/Codex/2026-09-10/diao/work/repos/spine-animation-ai") BACKGROUND = (26, 26, 48, 255) SCORE_WEIGHTS = { "foreground_mae_normalized": 0.45, "missing_foreground_ratio": 0.50, "extra_foreground_ratio": 0.25, "one_minus_mean_match": 0.20, "overlap_penalty": 0.30, "whole_character_risk": 0.80, "unmatched_part_ratio": 0.10, } def sha256(file: Path) -> str: return hashlib.sha256(file.read_bytes()).hexdigest() def finite(value, default=0.0) -> float: try: number = float(value) except (TypeError, ValueError): return default return number if math.isfinite(number) else default def save_json(file: Path, data): file.parent.mkdir(parents=True, exist_ok=True) file.write_text(json.dumps(data, ensure_ascii=False, indent=2, allow_nan=False) + "\n") def foreground_mask(image: Image.Image, threshold=240) -> np.ndarray: """Exclude only transparent/border-connected white background. This follows the corrected segment.py behavior and retains enclosed whites such as eye highlights, teeth and pale clothes. """ pixels = np.asarray(image.convert("RGBA")) transparent = pixels[:, :, 3] <= 128 white = np.all(pixels[:, :, :3] >= threshold, axis=2) _, labels = cv2.connectedComponents((white | transparent).astype(np.uint8), 8) border = np.unique(np.concatenate((labels[0], labels[-1], labels[:, 0], labels[:, -1]))) border = border[border != 0] return ~(np.isin(labels, border) | transparent) def bbox(mask: np.ndarray): ys, xs = np.nonzero(mask) if not len(xs): return (0, 0, 1, 1) return (int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1) def positioned_masks(parts_dir: Path, layout: dict, size: tuple[int, int]): """Use the same resize/rotate/center convention as compose_preview.py.""" masks = {} for name, pos in layout.get("parts", {}).items(): file = parts_dir / f"{name}.png" if not file.exists(): continue image = Image.open(file).convert("RGBA") width = max(1, round(finite(pos.get("width"), image.width))) height = max(1, round(finite(pos.get("height"), image.height))) image = image.resize((width, height), Image.Resampling.LANCZOS) angle = finite(pos.get("rotation")) if abs(angle) > 0.01: image = image.rotate(-angle, expand=True, resample=Image.Resampling.BICUBIC) cx = finite(pos.get("x")) + width / 2 cy = finite(pos.get("y")) + height / 2 layer = Image.new("L", size, 0) layer.paste(image.getchannel("A"), (round(cx - image.width / 2), round(cy - image.height / 2))) masks[name] = np.asarray(layer) > 128 return masks def evaluate(reference: Path, parts_dir: Path, layout_path: Path, assembled: Path) -> dict: ref = Image.open(reference).convert("RGBA") ref_mask = foreground_mask(ref) ref_pixels = max(1, int(ref_mask.sum())) layout = json.loads(layout_path.read_text()) stats = json.loads((parts_dir / "parts-manifest.json").read_text()) masks = positioned_masks(parts_dir, layout, ref.size) hits = np.zeros((ref.height, ref.width), dtype=np.uint16) for mask in masks.values(): hits += mask.astype(np.uint16) predicted_mask = hits > 0 predicted_pixels = max(1, int(predicted_mask.sum())) intersection = int((predicted_mask & ref_mask).sum()) union = max(1, int((predicted_mask | ref_mask).sum())) overlap_ratio = float((hits > 1).sum()) / predicted_pixels excess_coverage = float(np.maximum(hits.astype(np.int32) - 1, 0).sum()) / ref_pixels # Some overlap at joints or face/hair layers is expected. Penalize the # redundant area above a 15% allowance, while keeping raw values visible. overlap_penalty = min(2.0, max(0.0, excess_coverage - 0.15)) ref_display = Image.new("RGBA", ref.size, BACKGROUND) ref_display.alpha_composite(ref) ref_rgb = np.asarray(ref_display)[:, :, :3].astype(np.float32) out_rgb = np.asarray(Image.open(assembled).convert("RGB")).astype(np.float32) error = np.abs(out_rgb - ref_rgb) foreground_mae = float(error[ref_mask].mean()) if ref_mask.any() else 255.0 union_mae = float(error[predicted_mask | ref_mask].mean()) if union else 255.0 parts = layout.get("parts", {}) scores = [min(1.0, max(0.0, finite(part.get("score")))) for part in parts.values()] mean_match = float(np.mean(scores)) if scores else 0.0 sift = sum(part.get("method") == "sift" for part in parts.values()) component_count = len(stats.get("parts", [])) unmatched_ratio = max(0, component_count - len(parts)) / max(1, component_count) rx1, ry1, rx2, ry2 = bbox(ref_mask) ref_width, ref_height = max(1, rx2 - rx1), max(1, ry2 - ry1) sheet_areas = {part["name"]: int(part.get("area", 0)) for part in stats.get("parts", [])} sheet_total = max(1, sum(sheet_areas.values())) large_components = [] for name, mask in masks.items(): x1, y1, x2, y2 = bbox(mask) fg_coverage = float((mask & ref_mask).sum()) / ref_pixels height_ratio = (y2 - y1) / ref_height width_ratio = (x2 - x1) / ref_width # An almost-complete copy of the character gives deceptively good # SIFT/MAE, but is not useful as a rig layer. Require both substantial # foreground coverage and nearly full body extent to raise this risk. coverage_risk = float(np.clip((fg_coverage - 0.45) / 0.35, 0, 1)) height_risk = float(np.clip((height_ratio - 0.60) / 0.30, 0, 1)) width_risk = float(np.clip((width_ratio - 0.45) / 0.40, 0, 1)) risk = coverage_risk * height_risk * width_risk large_components.append({ "name": name, "risk": round(risk, 5), "reference_foreground_coverage": round(fg_coverage, 5), "reference_height_ratio": round(height_ratio, 5), "reference_width_ratio": round(width_ratio, 5), "sheet_foreground_share": round(sheet_areas.get(name, 0) / sheet_total, 5), }) large_components.sort(key=lambda item: item["risk"], reverse=True) whole_risk = max((item["risk"] for item in large_components), default=0.0) missing_ratio = max(0.0, 1 - intersection / ref_pixels) extra_ratio = float((predicted_mask & ~ref_mask).sum()) / ref_pixels terms = { "foreground_mae_normalized": foreground_mae / 255, "missing_foreground_ratio": missing_ratio, "extra_foreground_ratio": extra_ratio, "one_minus_mean_match": 1 - mean_match, "overlap_penalty": overlap_penalty, "whole_character_risk": whole_risk, "unmatched_part_ratio": unmatched_ratio, } weighted_terms = {name: value * SCORE_WEIGHTS[name] for name, value in terms.items()} return { "segmented_parts": component_count, "matched_parts": len(parts), "mean_match": round(mean_match, 5), "sift_count": sift, "template_count": sum(part.get("method") == "template" for part in parts.values()), "sift_inliers": sum(int(finite(part.get("n_matches"))) for part in parts.values()), "foreground_mae": round(foreground_mae, 4), "foreground_union_mae": round(union_mae, 4), "foreground_iou": round(intersection / union, 5), "missing_foreground_ratio": round(missing_ratio, 5), "extra_foreground_ratio": round(extra_ratio, 5), "overlap_ratio": round(overlap_ratio, 5), "excess_coverage_ratio": round(excess_coverage, 5), "overlap_penalty": round(overlap_penalty, 5), "whole_character_risk": round(whole_risk, 5), "large_components": large_components, "selection_score": round(sum(weighted_terms.values()), 6), "selection_score_terms": {key: round(value, 6) for key, value in weighted_terms.items()}, "manual_review_required": True, } def evaluate_candidate(task: dict, args) -> dict: candidate = task["candidate"] reference = task["reference"] sheet = candidate / "generated-parts-sheet.png" output = candidate / "resegmented" parts = output / "parts" layout = output / "layout.json" assembled = output / "assembled.png" report_file = output / "evaluation.json" output.mkdir(parents=True, exist_ok=True) script_files = [ROOT / "segment.py", ROOT / "compose_preview.py", args.genielabs_dir / "scripts" / "position_parts.py", Path(__file__)] fingerprint = { "reference_sha256": sha256(reference), "sheet_sha256": sha256(sheet), "script_sha256": {file.name: sha256(file) for file in script_files}, "min_area": args.min_area, "padding": args.padding, "bg_threshold": args.bg_threshold, } if args.resume and report_file.exists() and assembled.exists() and layout.exists(): old = json.loads(report_file.read_text()) if old.get("fingerprint") == fingerprint and old.get("status") == "completed": return old # Only remove files produced by this evaluator. Original parts and # generated sheets are inputs and never enter this cleanup path. if parts.exists(): shutil.rmtree(parts) parts.mkdir(parents=True) result = { "batch": task["batch"], "candidate": candidate.name, "reference": str(reference), "generated_sheet": str(sheet), "output_dir": str(output), "parts_dir": str(parts), "layout": str(layout), "assembled": str(assembled), "fingerprint": fingerprint, } commands = [ [sys.executable, str(ROOT / "segment.py"), "--atlas", str(sheet), "--output-dir", str(parts), "--min-area", str(args.min_area), "--padding", str(args.padding), "--bg-threshold", str(args.bg_threshold)], [sys.executable, str(args.genielabs_dir / "scripts" / "position_parts.py"), "--reference", str(reference), "--parts", str(parts), "--output", str(layout)], [sys.executable, str(ROOT / "compose_preview.py"), "--reference", str(reference), "--parts", str(parts), "--layout", str(layout), "--output", str(assembled)], ] try: with (output / "run.log").open("w") as log: for command in commands: log.write(" ".join(command) + "\n") log.flush() subprocess.run(command, check=True, stdout=log, stderr=subprocess.STDOUT, timeout=args.timeout) result.update(status="completed", metrics=evaluate(reference, parts, layout, assembled)) except Exception as error: result.update(status="failed", error=str(error)) save_json(report_file, result) return result def font(size=20): for file in ("/System/Library/Fonts/Supplemental/Arial.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"): if Path(file).exists(): return ImageFont.truetype(file, size) return ImageFont.load_default() def paste_panel(canvas: Image.Image, file: Path, box): x, y, width, height = box panel = Image.new("RGBA", (width, height), BACKGROUND) image = Image.open(file).convert("RGBA") image.thumbnail((width - 12, height - 12), Image.Resampling.LANCZOS) panel.alpha_composite(image, ((width - image.width) // 2, (height - image.height) // 2)) canvas.paste(panel.convert("RGB"), (x, y)) def contact_sheet(file: Path, title: str, rows: list[dict], selected: str | None): width, row_height, header, footer = 1320, 466, 85, 62 canvas = Image.new("RGB", (width, header + row_height * len(rows) + footer), (239, 241, 245)) draw = ImageDraw.Draw(canvas) draw.text((20, 13), title, font=font(25), fill=(25, 32, 45)) draw.text((20, 48), "Reference | Generated sheet | Re-segmented assembly (all matched parts)", font=font(19), fill=(60, 65, 80)) for index, result in enumerate(rows): y = header + index * row_height label = result["candidate"] if selected == result["candidate"]: label += " [lowest heuristic score]" draw.text((20, y + 3), label, font=font(21), fill=(20, 58, 100)) if result["status"] == "completed": metrics = result["metrics"] line = (f"parts {metrics['matched_parts']}/{metrics['segmented_parts']} SIFT {metrics['sift_count']} match {metrics['mean_match']:.3f} " f"MAE {metrics['foreground_mae']:.2f} overlap {metrics['overlap_penalty']:.2f} whole-body risk {metrics['whole_character_risk']:.2f} score {metrics['selection_score']:.3f}") draw.text((20, y + 30), line, font=font(17), fill=(40, 45, 55)) for image, box in ((result["reference"], (20, y + 62, 350, 388)), (result["generated_sheet"], (385, y + 62, 550, 388)), (result["assembled"], (950, y + 62, 350, 388))): paste_panel(canvas, Path(image), box) else: draw.text((20, y + 45), "FAILED: " + result.get("error", "unknown")[:130], font=font(17), fill=(160, 40, 40)) draw.text((20, canvas.height - 44), "Scores are diagnostic heuristics, not an anatomy or animation-quality approval. Inspect each candidate visually.", font=font(17), fill=(65, 65, 80)) file.parent.mkdir(parents=True, exist_ok=True) canvas.save(file, quality=92) def candidate_order(result): found = re.search(r"(\d+)$", result["candidate"]) return int(found.group(1)) if found else 0 def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--root", type=Path, default=ROOT / "output" / "prompt-optimization") parser.add_argument("--genielabs-dir", type=Path, default=DEFAULT_GENIELABS) parser.add_argument("--batch", action="append", help="Evaluate selected batch(es), default: all seven original batches") parser.add_argument("--jobs", type=int, default=2) parser.add_argument("--min-area", type=int, default=500) parser.add_argument("--padding", type=int, default=12) parser.add_argument("--bg-threshold", type=int, default=240) parser.add_argument("--timeout", type=int, default=600, help="Per-stage timeout in seconds") parser.add_argument("--resume", action="store_true", help="Reuse only evaluations with matching input/script fingerprints") args = parser.parse_args() args.root = args.root.resolve() args.genielabs_dir = args.genielabs_dir.resolve() if not (args.genielabs_dir / "scripts" / "position_parts.py").exists(): parser.error("Genielabs position_parts.py not found") tasks, missing, missing_sheets = [], [], [] for name in args.batch or DEFAULT_BATCHES: batch = args.root / name reference = next((batch / file for file in ("reference-position.png", "source.png", "source-original.png") if (batch / file).exists()), None) candidates = sorted((batch / "candidates").glob("attempt-*/generated-parts-sheet.png")) missing_sheets.extend({"batch": name, "candidate": directory.name, "reason": "no generated-parts-sheet.png; no offline input to evaluate"} for directory in sorted((batch / "candidates").glob("attempt-*")) if directory.is_dir() and not (directory / "generated-parts-sheet.png").exists()) if reference is None or not candidates: missing.append({"batch": name, "reason": "missing reference or candidate sheets"}) continue tasks.extend({"batch": name, "reference": reference, "candidate": sheet.parent} for sheet in candidates) results = [] print(f"Offline evaluation: {len(tasks)} existing candidates; no generation API calls.", flush=True) with ThreadPoolExecutor(max_workers=max(1, args.jobs)) as pool: futures = [pool.submit(evaluate_candidate, task, args) for task in tasks] for future in as_completed(futures): result = future.result() results.append(result) suffix = f"score={result['metrics']['selection_score']:.3f}" if result["status"] == "completed" else result.get("error", "failed") print(f"[{len(results)}/{len(tasks)}] {result['batch']}/{result['candidate']}: {suffix}", flush=True) batch_reports = [] for name in args.batch or DEFAULT_BATCHES: rows = sorted((result for result in results if result["batch"] == name), key=candidate_order) if not rows: continue completed = [result for result in rows if result["status"] == "completed"] ranked = sorted(completed, key=lambda result: result["metrics"]["selection_score"]) selected = ranked[0]["candidate"] if ranked else None output = args.root / name / "resegmented" report = { "batch": name, "offline": True, "original_artifacts_modified": False, "selected_by_heuristic": selected, "ranking": [result["candidate"] for result in ranked], "score_direction": "lower is better", "score_weights": SCORE_WEIGHTS, "manual_review_required": True, "candidates": rows, "missing_candidate_sheets": [item for item in missing_sheets if item["batch"] == name], "contact_sheet": str(output / "contact-sheet.jpg"), } save_json(output / "report.json", report) contact_sheet(output / "contact-sheet.jpg", f"{name}: offline candidate re-evaluation", rows, selected) batch_reports.append(report) summary = { "created_at": datetime.now(timezone.utc).isoformat(), "offline": True, "candidate_count": len(results), "completed": sum(result["status"] == "completed" for result in results), "failed": sum(result["status"] != "completed" for result in results), "score_weights": SCORE_WEIGHTS, "missing_batches": missing, "missing_candidate_sheets": missing_sheets, "batches": batch_reports, "limitations": [ "Matching scores from SIFT and template fallback are not calibrated probabilities.", "Layer overlap is sometimes intentional; raw overlap and penalized excess are both reported.", "A large complete-character component can produce low pixel error but is not a valid articulated decomposition; the risk term flags this shortcut.", "This evaluation measures static reconstruction with the existing compose_preview transform convention, not joint placement or animation quality.", ], } save_json(args.root / "reselection-report.json", summary) best_rows = [] for report in batch_reports: for row in report["candidates"]: if row["candidate"] == report["selected_by_heuristic"]: best_rows.append({**row, "candidate": f"{report['batch']} / {row['candidate']}"}) if best_rows: contact_sheet(args.root / "reselection-contact-sheet.jpg", "Offline re-evaluation: lowest heuristic score per character", best_rows, None) print(f"Report: {args.root / 'reselection-report.json'}", flush=True) return 1 if missing or any(result["status"] != "completed" for result in results) else 0 if __name__ == "__main__": raise SystemExit(main())