"""Run whole-image GPT Image post-processing over variation_skill_runs outputs.""" from __future__ import annotations import argparse import csv import html import os import re import sys from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from modules.full_image_polisher.full_image_polisher import ( DEFAULT_PROMPT, GLOBAL_COHERENCE_PROMPT, polish_full_image, polish_two_pass, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Post-process rendered variation PNGs with an image edit backend.") parser.add_argument("--runs-root", default="data/output/variation_skill_runs") parser.add_argument("--variation", default="", help="Comma-separated variation filter.") parser.add_argument("--run-dir", action="append", default=[], help="Specific run directory to process; can be repeated.") parser.add_argument( "--render-dir", default="latest_all_samples", help="Render directory name inside each run, or latest_all_samples for the highest render_all_samples_roundN.", ) parser.add_argument("--sample", default="sample_00", help="Sample filter, comma-separated, or all.") parser.add_argument("--output-subdir", default="postprocess_gpt_image2") parser.add_argument("--limit", type=int, default=None) parser.add_argument("--resume", action="store_true") parser.add_argument("--include-failed", action="store_true", help="Also process rows whose render ok field is not True.") parser.add_argument("--mode", choices=("global", "slot", "two_pass"), default="two_pass") parser.add_argument("--image-backend", choices=("auto", "openai", "pinco"), default="auto") parser.add_argument("--model", default="gpt-image-2") parser.add_argument("--size", default="auto") parser.add_argument("--quality", default="auto") parser.add_argument("--input-fidelity", choices=("high", "low"), default="high") parser.add_argument("--output-format", default="png") parser.add_argument("--prompt", default=DEFAULT_PROMPT) parser.add_argument("--data-context", choices=("none", "compact", "strict"), default="compact") parser.add_argument("--data-context-max-rows", type=int, default=80) parser.add_argument("--api-key-env", default="OPENAI_API_KEY") parser.add_argument("--base-url", default=None) parser.add_argument("--dry-run", action="store_true") parser.add_argument("--resize-to-input", action="store_true") parser.add_argument("--slot-guided", action="store_true", help="Use each final SVG to build slot mask/overlay guidance.") parser.add_argument("--slot-overlay-input", action="store_true", help="Send labeled slot overlay as a second image input. Off by default to avoid label leakage.") parser.add_argument("--include-title-slots", action="store_true") parser.add_argument("--slot-padding-px", type=int, default=8) parser.add_argument( "--pinco-command", default=None, help=( "Command template for local Pinco inference. Placeholders: {input}, {mask}, " "{foreground}, {prompt_file}, {output}, {model}, {width}, {height}." ), ) parser.add_argument("--pinco-url", default=None, help="HTTP endpoint for a Pinco inpainting service.") parser.add_argument("--pinco-timeout", type=int, default=600) return parser.parse_args() def _variation_filter(spec: str) -> set[str]: return {item.strip() for item in spec.split(",") if item.strip()} def discover_run_dirs(args: argparse.Namespace) -> list[Path]: if args.run_dir: return [Path(path) for path in args.run_dir] runs_root = Path(args.runs_root) wanted = _variation_filter(args.variation) run_dirs: list[Path] = [] for variation_dir in sorted(path for path in runs_root.iterdir() if path.is_dir()): if wanted and variation_dir.name not in wanted: continue dated = sorted( [path for path in variation_dir.iterdir() if path.is_dir()], key=lambda path: path.stat().st_mtime, reverse=True, ) if dated: run_dirs.append(dated[0]) return run_dirs def _resolve_row_path(path_text: str) -> Path: path = Path(path_text) if path.is_absolute(): return path root_path = (ROOT / path).resolve() if root_path.exists(): return root_path return (ROOT.parent / path).resolve() def _round_number(path: Path) -> int: match = re.search(r"round(\d+)$", path.name) return int(match.group(1)) if match else -1 def resolve_render_dir(run_dir: Path, render_dir_name: str) -> Path | None: if render_dir_name != "latest_all_samples": candidate = run_dir / render_dir_name return candidate if candidate.is_dir() else None candidates = [path for path in run_dir.glob("render_all_samples_round*") if path.is_dir()] if not candidates: return None return sorted(candidates, key=lambda path: (_round_number(path), path.stat().st_mtime), reverse=True)[0] def sample_filter(spec: str) -> set[str] | None: if spec == "all": return None return {item.strip() for item in spec.split(",") if item.strip()} def _png_for_row(row: dict[str, str]) -> Path | None: final_svg = row.get("final_svg", "") if not final_svg: return None svg_path = _resolve_row_path(final_svg) png_path = svg_path.with_suffix(".png") if png_path.is_file(): return png_path candidates = sorted(svg_path.parent.glob(f"{svg_path.stem}*.png")) return candidates[-1] if candidates else None def _svg_for_row(row: dict[str, str]) -> Path | None: final_svg = row.get("final_svg", "") if not final_svg: return None svg_path = _resolve_row_path(final_svg) return svg_path if svg_path.is_file() else None def _data_for_row(row: dict[str, str]) -> Path | None: final_svg = row.get("final_svg", "") if not final_svg: return None svg_path = _resolve_row_path(final_svg) candidates = [ svg_path.with_suffix("") / "data.json", svg_path.parent / "data.json", ] chart_svg = row.get("chart_svg", "") if chart_svg: chart_svg_path = _resolve_row_path(chart_svg) candidates.append(chart_svg_path.parent / "data.json") for candidate in candidates: if candidate.is_file(): return candidate return None def _reference_map_for_run(run_dir: Path) -> Path | None: candidate = run_dir / "reference_element_map.json" return candidate if candidate.is_file() else None def load_jobs(run_dirs: list[Path], render_dir_name: str, sample_spec: str, include_failed: bool) -> list[dict[str, str]]: jobs: list[dict[str, str]] = [] wanted_samples = sample_filter(sample_spec) for run_dir in run_dirs: render_dir = resolve_render_dir(run_dir, render_dir_name) if render_dir is None: continue tasks_csv = render_dir / "_tasks.csv" if not tasks_csv.is_file(): continue with tasks_csv.open(newline="", encoding="utf-8") as fh: for row in csv.DictReader(fh): if wanted_samples is not None and row.get("sample") not in wanted_samples: continue if not include_failed and row.get("ok") != "True": continue png_path = _png_for_row(row) if not png_path: continue svg_path = _svg_for_row(row) data_path = _data_for_row(row) jobs.append({ "run_dir": str(run_dir), "render_dir": render_dir.name, "chart_name": row.get("chart_name", ""), "sample": row.get("sample", ""), "input_png": str(png_path), "input_svg": str(svg_path) if svg_path else "", "input_data": str(data_path) if data_path else "", "reference_map": str(_reference_map_for_run(run_dir) or ""), }) return jobs def read_done(tasks_csv: Path) -> set[str]: if not tasks_csv.is_file(): return set() with tasks_csv.open(newline="", encoding="utf-8") as fh: return {row["input_png"] for row in csv.DictReader(fh) if row.get("ok") == "True"} def output_for(job: dict[str, str], output_subdir: str, mode: str = "global") -> Path: run_dir = Path(job["run_dir"]) input_png = Path(job["input_png"]) suffix_by_mode = { "global": "coherence_polished", "slot": "slot_guided_polished", "two_pass": "two_pass_final", } suffix = suffix_by_mode[mode] rel_parts = [job["render_dir"], job["chart_name"], f"{input_png.stem}.{suffix}.png"] return run_dir / output_subdir / Path(*rel_parts) def write_preview(preview_path: Path, rows: list[dict[str, str]]) -> None: preview_path.parent.mkdir(parents=True, exist_ok=True) parent = preview_path.parent.resolve() def rel(path_text: str) -> str: if not path_text: return "" path = Path(path_text) if not path.is_absolute(): path = (ROOT / path).resolve() return os.path.relpath(path, parent) parts = [ "", "", "

GPT Image Whole-Image Postprocess

", f"

tasks={len(rows)} ok={sum(row['ok']=='True' for row in rows)} fail={sum(row['ok']!='True' for row in rows)}

", "
", ] for row in rows: cls = "card" if row["ok"] == "True" else "card bad" source = html.escape(rel(row["input_png"])) output = html.escape(rel(row["output_png"])) output_img = f"" if row["ok"] == "True" and output else "
no output
" parts.append( f"
{output_img}
" f"
{html.escape(row['chart_name'])} {html.escape(row['sample'])}
" f"ok={row['ok']} {html.escape(row.get('err', ''))}
{html.escape(row['output_png'])}
" ) parts.append("
") preview_path.write_text("\n".join(parts), encoding="utf-8") def main() -> int: args = parse_args() try: from config import api_key, base_url except Exception: api_key = None base_url = None run_dirs = discover_run_dirs(args) jobs = load_jobs(run_dirs, args.render_dir, args.sample, args.include_failed) if args.limit is not None: jobs = jobs[: args.limit] output_roots = {Path(job["run_dir"]) / args.output_subdir for job in jobs} for output_root in output_roots: output_root.mkdir(parents=True, exist_ok=True) task_logs = {output_root: output_root / "postprocess_tasks.csv" for output_root in output_roots} done: set[str] = set() if args.resume: for tasks_csv in task_logs.values(): done.update(read_done(tasks_csv)) jobs = [job for job in jobs if job["input_png"] not in done] fieldnames = [ "run_dir", "render_dir", "chart_name", "sample", "input_png", "input_svg", "input_data", "reference_map", "output_png", "manifest", "ok", "err", ] rows_by_output_root: dict[Path, list[dict[str, str]]] = {root: [] for root in output_roots} effective_mode = "slot" if args.slot_guided else args.mode for index, job in enumerate(jobs, 1): output_png = output_for(job, args.output_subdir, mode=effective_mode) output_root = Path(job["run_dir"]) / args.output_subdir row = {**job, "output_png": str(output_png), "manifest": "", "ok": "False", "err": ""} try: if effective_mode == "two_pass": if not job.get("input_svg"): raise ValueError("two_pass mode requires input_svg from _tasks.csv") manifest = polish_two_pass( input_png=Path(job["input_png"]), output_png=output_png, api_key=api_key, base_url=args.base_url or base_url, svg_path=Path(job["input_svg"]), reference_map=Path(job["reference_map"]) if job.get("reference_map") else None, data_json=Path(job["input_data"]) if job.get("input_data") else None, data_context_mode=args.data_context, data_context_max_rows=args.data_context_max_rows, model=args.model, size=args.size, quality=args.quality, input_fidelity=args.input_fidelity, output_format=args.output_format, api_key_env=args.api_key_env, dry_run=args.dry_run, resize_to_input=args.resize_to_input, slot_overlay_input=args.slot_overlay_input, include_title_slots=args.include_title_slots, slot_padding_px=args.slot_padding_px, image_backend=args.image_backend, pinco_command=args.pinco_command, pinco_url=args.pinco_url, pinco_timeout=args.pinco_timeout, ) else: mode_prompt = args.prompt if effective_mode == "global" and mode_prompt == DEFAULT_PROMPT: mode_prompt = GLOBAL_COHERENCE_PROMPT manifest = polish_full_image( input_png=Path(job["input_png"]), output_png=output_png, api_key=api_key, base_url=args.base_url or base_url, svg_path=Path(job["input_svg"]) if effective_mode == "slot" and job.get("input_svg") else None, reference_map=Path(job["reference_map"]) if job.get("reference_map") else None, data_json=Path(job["input_data"]) if job.get("input_data") else None, data_context_mode=args.data_context if effective_mode == "global" else "none", data_context_max_rows=args.data_context_max_rows, model=args.model, size=args.size, quality=args.quality, input_fidelity=args.input_fidelity, output_format=args.output_format, prompt=mode_prompt, api_key_env=args.api_key_env, dry_run=args.dry_run, resize_to_input=args.resize_to_input, slot_guided=effective_mode == "slot", slot_overlay_input=args.slot_overlay_input, include_title_slots=args.include_title_slots, slot_padding_px=args.slot_padding_px, image_backend=args.image_backend, pinco_command=args.pinco_command, pinco_url=args.pinco_url, pinco_timeout=args.pinco_timeout, ) row["manifest"] = str(Path(manifest["output_png"]).with_name(f"{Path(manifest['output_png']).stem}.manifest.json")) row["ok"] = "True" except Exception as exc: row["err"] = f"{type(exc).__name__}: {exc}"[:500] rows_by_output_root.setdefault(output_root, []).append(row) print(f"[{index}/{len(jobs)}] {job['chart_name']}/{job['sample']} ok={row['ok']} output={row['output_png']}") for output_root, rows in rows_by_output_root.items(): if not rows: continue tasks_csv = output_root / "postprocess_tasks.csv" mode = "a" if args.resume and tasks_csv.exists() else "w" with tasks_csv.open(mode, newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=fieldnames) if mode == "w": writer.writeheader() writer.writerows(rows) all_rows = list(csv.DictReader(tasks_csv.open(newline="", encoding="utf-8"))) write_preview(output_root / "preview.html", all_rows) print(f"Wrote {tasks_csv}") print(f"Wrote {output_root / 'preview.html'}") return 0 if all(row["ok"] == "True" for rows in rows_by_output_root.values() for row in rows) else 1 if __name__ == "__main__": raise SystemExit(main())