Buckets:
| """Runtime orchestration for offline Visual-Jigsaw visual-quality scoring.""" | |
| from __future__ import annotations | |
| import argparse | |
| import gc | |
| import hashlib | |
| import json | |
| import os | |
| from collections import defaultdict | |
| from pathlib import Path | |
| import torch | |
| from PIL import Image | |
| from sft.runtime import acquire_run_writer, claim_run_directory | |
| from .full_data import load_item_images | |
| from .generated_codes import load_generated_codes | |
| from .visual_quality import ( | |
| DEFAULT_ASPECT_TOLERANCE, | |
| DEFAULT_BLANK_STD_THRESHOLD, | |
| DEFAULT_DUPLICATE_DINO_THRESHOLD, | |
| DEFAULT_DUPLICATE_LPIPS_THRESHOLD, | |
| SAMPLE_SCHEMA, | |
| atomic_write_json, | |
| blank_image_metrics, | |
| generated_state_events, | |
| geometry_metrics, | |
| native_code_metrics, | |
| prediction_rows, | |
| read_json, | |
| sample_artifact_dirs, | |
| summarize_samples, | |
| target_code_images, | |
| PackedItemCache, | |
| ) | |
| from .visual_quality_models import RolloutRGBDecoder, VisualMetricModels | |
| DINO_REVISION = "f9e44c814b77203eaa57a6bdbbd535f21ede1415" | |
| CLIP_REVISION = "3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268" | |
| def _seed(sample_id: str, image_index: int) -> int: | |
| payload = hashlib.sha256( | |
| f"20260827:{sample_id}:{image_index}".encode() | |
| ).digest()[:8] | |
| return int.from_bytes(payload, "big") & ((1 << 63) - 1) | |
| def _valid_png(path: Path) -> bool: | |
| if not path.is_file(): | |
| return False | |
| try: | |
| with Image.open(path) as image: | |
| image.verify() | |
| except (OSError, ValueError): | |
| return False | |
| return True | |
| def _atomic_save_png(image: Image.Image, path: Path) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") | |
| try: | |
| image.save(temporary, format="PNG") | |
| os.replace(temporary, path) | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| def _request(args: argparse.Namespace, rows: list[dict]) -> dict: | |
| eval_arguments = read_json(args.eval_run / "arguments.json") | |
| return { | |
| "schema": "zebra_visual_quality_request_v1", | |
| "eval_run": str(args.eval_run.resolve()), | |
| "eval_model_label": eval_arguments["model_label"], | |
| "cache_root": str(args.cache_root.resolve()), | |
| "model_path": str(args.model_path.resolve()), | |
| "source_root": str(args.source_root.resolve()), | |
| "graph_root": str(args.graph_root.resolve()), | |
| "samples": len(rows), | |
| "sample_ids_sha256": hashlib.sha256( | |
| "\n".join(row["sample_id"] for row in rows).encode() | |
| ).hexdigest(), | |
| "source_snapshot_sha256": os.environ.get("SPATIAL_MMCOT_SOURCE_HASH"), | |
| "decoder": { | |
| "steps": args.decoder_steps, | |
| "guidance": args.decoder_guidance, | |
| "upsampling": args.decoder_upsampling, | |
| "batch_size": args.decoder_batch_size, | |
| "seed_contract": "sha256('20260827:{sample_id}:{one_based_index}')", | |
| }, | |
| "metrics": { | |
| "piq": "0.8.0", | |
| "dino_path": str(args.dino_path.resolve()), | |
| "dino_revision": DINO_REVISION, | |
| "clip_path": str(args.clip_path.resolve()), | |
| "clip_revision": CLIP_REVISION, | |
| "semantic_batch_size": args.semantic_batch_size, | |
| "aspect_tolerance": args.aspect_tolerance, | |
| "blank_std_threshold": args.blank_std_threshold, | |
| "duplicate_dino_threshold": args.duplicate_dino_threshold, | |
| "duplicate_lpips_threshold": args.duplicate_lpips_threshold, | |
| }, | |
| } | |
| def _prepare_output(args: argparse.Namespace, rows: list[dict]) -> tuple[Path, object]: | |
| output_dir = claim_run_directory(args.output_dir, resume=args.resume) | |
| writer = acquire_run_writer(output_dir) | |
| request = _request(args, rows) | |
| request_path = output_dir / "arguments.json" | |
| if request_path.exists(): | |
| if read_json(request_path) != request: | |
| writer.close() | |
| raise ValueError("visual-quality resume arguments differ") | |
| else: | |
| atomic_write_json(request_path, request) | |
| return output_dir, writer | |
| def _decode_sample( | |
| decoder: RolloutRGBDecoder, | |
| *, | |
| row: dict, | |
| codes_path: Path, | |
| output_dir: Path, | |
| args: argparse.Namespace, | |
| ) -> None: | |
| sample_id = row["sample_id"] | |
| sample_dir = output_dir / "samples" / sample_id | |
| status_path = sample_dir / "decode.json" | |
| expected_paths = [ | |
| sample_dir / f"generated_{index:02d}.png" | |
| for index in range(1, int(row["generated_images"]) + 1) | |
| ] | |
| if status_path.exists() and all(_valid_png(path) for path in expected_paths): | |
| status = read_json(status_path) | |
| if status.get("generated_images") == row["generated_images"]: | |
| return | |
| codes = load_generated_codes(codes_path) | |
| grids = [tuple(int(value) for value in grid) for grid in row["generated_grids"]] | |
| if len(codes) != len(grids) or len(codes) != int(row["generated_images"]): | |
| raise ValueError(f"generated code/grid count mismatch for {sample_id}") | |
| by_grid: dict[tuple[int, int], list[tuple[int, torch.Tensor]]] = defaultdict(list) | |
| for index, (code, grid) in enumerate(zip(codes, grids, strict=True), start=1): | |
| if code is None: | |
| raise ValueError(f"control rollout has no codes for {sample_id}:{index}") | |
| destination = sample_dir / f"generated_{index:02d}.png" | |
| if not _valid_png(destination): | |
| by_grid[grid].append((index, code)) | |
| for grid, pending in sorted(by_grid.items()): | |
| for start in range(0, len(pending), args.decoder_batch_size): | |
| chunk = pending[start : start + args.decoder_batch_size] | |
| generators = [ | |
| torch.Generator(device=decoder.device).manual_seed(_seed(sample_id, index)) | |
| for index, _code in chunk | |
| ] | |
| images = decoder.decode( | |
| [code for _index, code in chunk], | |
| grid=grid, | |
| generators=generators, | |
| steps=args.decoder_steps, | |
| guidance=args.decoder_guidance, | |
| upsampling=args.decoder_upsampling, | |
| ) | |
| if len(images) != len(chunk): | |
| raise RuntimeError("RGB decoder returned the wrong batch size") | |
| for (index, _code), image in zip(chunk, images, strict=True): | |
| _atomic_save_png( | |
| image, sample_dir / f"generated_{index:02d}.png" | |
| ) | |
| if not all(_valid_png(path) for path in expected_paths): | |
| raise RuntimeError(f"decoded PNG set is incomplete for {sample_id}") | |
| atomic_write_json( | |
| status_path, | |
| { | |
| "schema": "zebra_visual_quality_decode_v1", | |
| "sample_id": sample_id, | |
| "generated_images": len(codes), | |
| "generated_grids": [list(grid) for grid in grids], | |
| }, | |
| ) | |
| def decode_all( | |
| args: argparse.Namespace, | |
| rows: list[dict], | |
| artifacts: dict[str, Path], | |
| output_dir: Path, | |
| ) -> None: | |
| decoder = RolloutRGBDecoder(args.model_path, torch.device("cuda")) | |
| for position, row in enumerate(rows, start=1): | |
| _decode_sample( | |
| decoder, | |
| row=row, | |
| codes_path=artifacts[row["sample_id"]] / "generated_codes.pt", | |
| output_dir=output_dir, | |
| args=args, | |
| ) | |
| print( | |
| f"DECODE {position}/{len(rows)} sample={row['sample_id']} " | |
| f"images={row['generated_images']}", | |
| flush=True, | |
| ) | |
| del decoder | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| def _cosine(first: torch.Tensor, second: torch.Tensor) -> float: | |
| return float((first * second).sum()) | |
| def _score_sample( | |
| models: VisualMetricModels, | |
| *, | |
| row: dict, | |
| item: dict, | |
| artifact_dir: Path, | |
| output_dir: Path, | |
| args: argparse.Namespace, | |
| ) -> dict: | |
| sample_id = row["sample_id"] | |
| generated = [ | |
| Image.open(output_dir / "samples" / sample_id / f"generated_{index:02d}.png") | |
| .convert("RGB") | |
| .copy() | |
| for index in range(1, int(row["generated_images"]) + 1) | |
| ] | |
| generated_quality = [ | |
| { | |
| "image_index": index, | |
| "generated_grid": row["generated_grids"][index - 1], | |
| **blank_image_metrics( | |
| image, std_threshold=args.blank_std_threshold | |
| ), | |
| } | |
| for index, image in enumerate(generated, start=1) | |
| ] | |
| count_match = int(row["expected_images"]) == int(row["generated_images"]) | |
| events = generated_state_events(item) | |
| if len(events) != int(row["expected_images"]): | |
| raise ValueError(f"target event count mismatch for {sample_id}") | |
| references_by_key = ( | |
| load_item_images(item, args.source_root, args.graph_root) | |
| if count_match | |
| else {} | |
| ) | |
| references = [references_by_key[event["key"]] for event in events] if count_match else [] | |
| semantic_images = [*generated, *references] | |
| if semantic_images: | |
| dino, clip = models.embeddings( | |
| semantic_images, batch_size=args.semantic_batch_size | |
| ) | |
| else: | |
| dino = clip = torch.empty(0) | |
| generated_dino = dino[: len(generated)] | |
| generated_clip = clip[: len(generated)] | |
| reference_dino = dino[len(generated) :] | |
| reference_clip = clip[len(generated) :] | |
| target_codes = target_code_images(item) if count_match else [] | |
| rollout_codes = load_generated_codes(artifact_dir / "generated_codes.pt") | |
| paired = [] | |
| if count_match: | |
| for index, (image, reference, target, code, event) in enumerate( | |
| zip( | |
| generated, | |
| references, | |
| target_codes, | |
| rollout_codes, | |
| events, | |
| strict=True, | |
| ), | |
| start=1, | |
| ): | |
| geometry = geometry_metrics( | |
| reference.size, | |
| image.size, | |
| aspect_tolerance=args.aspect_tolerance, | |
| ) | |
| native = native_code_metrics(code, target) | |
| pixel_perceptual = ( | |
| models.aligned_metrics(image, reference) | |
| if geometry["geometry_comparable"] | |
| else None | |
| ) | |
| paired.append( | |
| { | |
| "image_index": index, | |
| "reference_key": event["key"], | |
| "expected_grid": row["expected_grids"][index - 1], | |
| "generated_grid": row["generated_grids"][index - 1], | |
| "geometry": geometry, | |
| "native": native, | |
| "pixel_perceptual": pixel_perceptual, | |
| "semantic": { | |
| "dinov2_cosine": _cosine( | |
| generated_dino[index - 1], reference_dino[index - 1] | |
| ), | |
| "clip_cosine": _cosine( | |
| generated_clip[index - 1], reference_clip[index - 1] | |
| ), | |
| }, | |
| } | |
| ) | |
| adjacent = [] | |
| for index in range(len(generated) - 1): | |
| geometry = geometry_metrics( | |
| generated[index].size, | |
| generated[index + 1].size, | |
| aspect_tolerance=args.aspect_tolerance, | |
| ) | |
| lpips = ( | |
| models.aligned_lpips(generated[index + 1], generated[index]) | |
| if geometry["geometry_comparable"] | |
| else None | |
| ) | |
| dino_cosine = _cosine(generated_dino[index], generated_dino[index + 1]) | |
| clip_cosine = _cosine(generated_clip[index], generated_clip[index + 1]) | |
| adjacent.append( | |
| { | |
| "first_image_index": index + 1, | |
| "second_image_index": index + 2, | |
| "dinov2_cosine": dino_cosine, | |
| "clip_cosine": clip_cosine, | |
| "lpips": lpips, | |
| "near_duplicate": bool( | |
| lpips is not None | |
| and dino_cosine >= args.duplicate_dino_threshold | |
| and lpips <= args.duplicate_lpips_threshold | |
| ), | |
| } | |
| ) | |
| return { | |
| "schema": SAMPLE_SCHEMA, | |
| "sample_id": sample_id, | |
| "evaluation_index": int(row["evaluation_index"]), | |
| "strict_correct": bool(row["strict_correct"]), | |
| "expected_images": int(row["expected_images"]), | |
| "generated_images": int(row["generated_images"]), | |
| "count_match": count_match, | |
| "exact_grid_match": row["expected_grids"] == row["generated_grids"], | |
| "generated_quality": generated_quality, | |
| "paired_images": paired, | |
| "adjacent_generated": adjacent, | |
| } | |
| def _write_report(path: Path, summary: dict) -> None: | |
| lines = [ | |
| "# Visual-Jigsaw Rollout Visual Quality", | |
| "", | |
| f"- Eval: `{summary['eval_run']}`", | |
| f"- Samples: {summary['samples']}", | |
| f"- Count matched: {summary['count_matched_samples']}", | |
| f"- Generated images: {summary['generated_images']}", | |
| f"- Paired images: {summary['paired_images']}", | |
| f"- Exact resolution pairs: {summary['exact_resolution_pairs']}", | |
| f"- Geometry-comparable pairs: {summary['geometry_comparable_pairs']}", | |
| f"- Blank fraction: {summary['blank_fraction']}", | |
| f"- Near-duplicate fraction: {summary['near_duplicate_fraction']}", | |
| "", | |
| "See `summary.json` and `sample_metrics.jsonl` for metric distributions.", | |
| ] | |
| path.write_text("\n".join(lines) + "\n", encoding="utf-8") | |
| def score_all( | |
| args: argparse.Namespace, | |
| rows: list[dict], | |
| artifacts: dict[str, Path], | |
| output_dir: Path, | |
| ) -> dict: | |
| models = VisualMetricModels( | |
| device=torch.device("cuda"), | |
| dino_path=args.dino_path, | |
| clip_path=args.clip_path, | |
| ) | |
| cache = PackedItemCache(args.cache_root) | |
| for position, row in enumerate(rows, start=1): | |
| sample_path = output_dir / "samples" / row["sample_id"] / "metrics.json" | |
| if sample_path.exists(): | |
| existing = read_json(sample_path) | |
| if existing.get("schema") == SAMPLE_SCHEMA: | |
| print( | |
| f"SCORE {position}/{len(rows)} existing={row['sample_id']}", | |
| flush=True, | |
| ) | |
| continue | |
| result = _score_sample( | |
| models, | |
| row=row, | |
| item=cache.item(row["sample_id"]), | |
| artifact_dir=artifacts[row["sample_id"]], | |
| output_dir=output_dir, | |
| args=args, | |
| ) | |
| atomic_write_json(sample_path, result) | |
| print(f"SCORE {position}/{len(rows)} sample={row['sample_id']}", flush=True) | |
| samples = [ | |
| read_json(output_dir / "samples" / row["sample_id"] / "metrics.json") | |
| for row in rows | |
| ] | |
| summary = summarize_samples(samples, eval_run=str(args.eval_run.resolve())) | |
| with (output_dir / "sample_metrics.jsonl").open("w", encoding="utf-8") as handle: | |
| for sample in samples: | |
| handle.write(json.dumps(sample, ensure_ascii=False, sort_keys=True) + "\n") | |
| atomic_write_json(output_dir / "summary.json", summary) | |
| _write_report(output_dir / "report.md", summary) | |
| (output_dir / "COMPLETE").write_text("complete\n", encoding="utf-8") | |
| return summary | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--eval-run", type=Path, required=True) | |
| parser.add_argument("--output-dir", type=Path, required=True) | |
| parser.add_argument("--cache-root", type=Path) | |
| parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT")) | |
| parser.add_argument( | |
| "--source-root", type=Path, default=Path("data/Zebra-CoT-Visual-Jigsaw") | |
| ) | |
| parser.add_argument( | |
| "--graph-root", type=Path, default=Path("data/Zebra-CoT-Visual-Jigsaw") | |
| ) | |
| parser.add_argument( | |
| "--dino-path", type=Path, default=Path("models/visual-quality/dinov2-base") | |
| ) | |
| parser.add_argument( | |
| "--clip-path", | |
| type=Path, | |
| default=Path("models/visual-quality/clip-vit-base-patch32"), | |
| ) | |
| parser.add_argument("--phase", choices=("all", "decode", "score"), default="all") | |
| parser.add_argument("--decoder-steps", type=int, default=28) | |
| parser.add_argument("--decoder-guidance", type=float, default=1.5) | |
| parser.add_argument("--decoder-upsampling", type=float, default=1.0) | |
| parser.add_argument("--decoder-batch-size", type=int, default=4) | |
| parser.add_argument("--semantic-batch-size", type=int, default=16) | |
| parser.add_argument("--aspect-tolerance", type=float, default=DEFAULT_ASPECT_TOLERANCE) | |
| parser.add_argument( | |
| "--blank-std-threshold", type=float, default=DEFAULT_BLANK_STD_THRESHOLD | |
| ) | |
| parser.add_argument( | |
| "--duplicate-dino-threshold", | |
| type=float, | |
| default=DEFAULT_DUPLICATE_DINO_THRESHOLD, | |
| ) | |
| parser.add_argument( | |
| "--duplicate-lpips-threshold", | |
| type=float, | |
| default=DEFAULT_DUPLICATE_LPIPS_THRESHOLD, | |
| ) | |
| parser.add_argument("--max-samples", type=int) | |
| parser.add_argument("--resume", action="store_true") | |
| return parser | |
| def main() -> None: | |
| args = build_parser().parse_args() | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("full visual-quality scoring requires a CUDA allocation") | |
| if args.decoder_batch_size <= 0 or args.semantic_batch_size <= 0: | |
| raise ValueError("batch sizes must be positive") | |
| if args.max_samples is not None and args.max_samples <= 0: | |
| raise ValueError("max_samples must be positive") | |
| args.eval_run = args.eval_run.resolve() | |
| eval_arguments = read_json(args.eval_run / "arguments.json") | |
| args.cache_root = ( | |
| args.cache_root.resolve() | |
| if args.cache_root is not None | |
| else Path(eval_arguments["cache_root"]).resolve() | |
| ) | |
| rows = prediction_rows(args.eval_run) | |
| if args.max_samples is not None: | |
| rows = rows[: args.max_samples] | |
| artifacts = sample_artifact_dirs(args.eval_run) | |
| output_dir, writer = _prepare_output(args, rows) | |
| try: | |
| if args.phase in {"all", "decode"}: | |
| decode_all(args, rows, artifacts, output_dir) | |
| if args.phase in {"all", "score"}: | |
| summary = score_all(args, rows, artifacts, output_dir) | |
| print("VISUAL_QUALITY_COMPLETE " + json.dumps(summary), flush=True) | |
| finally: | |
| writer.close() | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 18.4 kB
- Xet hash:
- 6266772e5267b20e10e6136347e6928f30b9313bc0e7546e7d0eb9336ad59168
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.