377 lines
19 KiB
Python
377 lines
19 KiB
Python
#!/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())
|