#!/usr/bin/env python from __future__ import annotations import argparse import hashlib import io import json import subprocess import sys from collections import Counter from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) import numpy as np # noqa: E402 from cil.chart_features import OBJECT_LAYOUT_EMBED_DIM # noqa: E402 EXTRACTOR_NAME = "rgb_object_layout_v1" def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description=( "Decode deployment-visible observation_ref JPEGs and write a " "deterministic RGB object-layout embedding into chart metadata." ) ) parser.add_argument( "--indexes", nargs="+", type=Path, default=[ Path("data/cil_charts_rgb_refs/train/index.json"), Path("data/cil_charts_rgb_refs/val/index.json"), Path("data/cil_charts_rgb_refs/test/index.json"), ], ) parser.add_argument( "--out-dir", type=Path, default=Path("runs/chart_object_embeddings_rgb_refs"), ) parser.add_argument("--overwrite", action="store_true") args = parser.parse_args(argv) out_dir = args.out_dir out_dir.mkdir(parents=True, exist_ok=True) _write_provenance(out_dir, args) split_rows = [] for index_path in args.indexes: split_rows.append(_process_index(index_path, overwrite=args.overwrite)) metrics = { "report_type": "chart_object_embedding_export", "schema_version": 1, "embedding_dim": OBJECT_LAYOUT_EMBED_DIM, "extractor": EXTRACTOR_NAME, "indexes": [str(path) for path in args.indexes], "splits": split_rows, "data_hash": {row["split"]: row["content_hash_after"] for row in split_rows}, "split_hash": {row["split"]: row["split_hash"] for row in split_rows}, "leakage_contract": { "reads_outcomes": False, "reads_observation_ref": True, "writes_object_embedding_path": True, }, } (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n") (out_dir / "metrics_by_task.json").write_text(_metrics_by_task(split_rows) + "\n") (out_dir / "metrics_by_seed.json").write_text("{}\n") (out_dir / "table.tex").write_text(_table(split_rows) + "\n") (out_dir / "report.md").write_text(_report(metrics) + "\n") (out_dir / "train.log").write_text("not a training run; exported object-layout embeddings\n") (out_dir / "eval.log").write_text( "decoded observation_ref JPEGs only; no outcomes, labels, or hidden branches read\n" ) print(json.dumps({"out_dir": str(out_dir), "splits": len(split_rows)}, indent=2)) return 0 def _process_index(index_path: Path, *, overwrite: bool) -> dict[str, Any]: index = json.loads(index_path.read_text()) split = str(index.get("split", index_path.parent.name)) embed_path = index_path.parent / "object_embeddings_rgb_layout.npz" if embed_path.exists() and not overwrite: raise FileExistsError(f"{embed_path} exists; pass --overwrite to replace it") ref_to_row: dict[tuple[str, str], int] = {} embeddings: list[np.ndarray] = [] counters: Counter[str] = Counter() task_counts: Counter[str] = Counter() updated_shards: list[str] = [] h5_cache: dict[Path, Any] = {} try: for shard in index.get("shards", []): shard_path = index_path.parent / str(shard["path"]) with np.load(shard_path, allow_pickle=False) as data: arrays = {key: data[key] for key in data.files} metadata_values = arrays["metadata_json"] updated_metadata = [] for raw in metadata_values: metadata = _json_loads(str(raw)) task_counts[str(metadata.get("task_id", "unknown"))] += 1 counters["rows"] += 1 source_dataset = str(metadata.get("source_dataset", "")) observation_ref = str(metadata.get("observation_ref", "")) if not source_dataset or not observation_ref: counters["missing_observation_ref"] += 1 updated_metadata.append(json.dumps(metadata, sort_keys=True)) continue key = (source_dataset, observation_ref) if key not in ref_to_row: ref_to_row[key] = len(embeddings) embeddings.append(_embedding_for_ref(key, h5_cache)) metadata["object_embedding_path"] = ( f"{embed_path.name}#embeddings/{ref_to_row[key]}" ) metadata["object_embedding_extractor"] = EXTRACTOR_NAME metadata["object_embedding_dim"] = OBJECT_LAYOUT_EMBED_DIM counters["rows_with_embedding"] += 1 updated_metadata.append(json.dumps(metadata, sort_keys=True)) arrays["metadata_json"] = np.asarray(updated_metadata) np.savez_compressed(shard_path, **arrays) updated_shards.append(str(shard_path)) finally: for handle in h5_cache.values(): handle.close() embedding_matrix = ( np.stack(embeddings).astype(np.float32) if embeddings else np.zeros((0, OBJECT_LAYOUT_EMBED_DIM), dtype=np.float32) ) np.savez_compressed( embed_path, embeddings=embedding_matrix, extractor=np.asarray([EXTRACTOR_NAME]), observation_refs=np.asarray([ref for _, ref in ref_to_row]), source_datasets=np.asarray([source for source, _ in ref_to_row]), ) index["object_embedding_manifest"] = { "path": embed_path.name, "dataset": "embeddings", "dim": OBJECT_LAYOUT_EMBED_DIM, "extractor": EXTRACTOR_NAME, "num_embeddings": int(embedding_matrix.shape[0]), "reads_outcomes": False, } index["shard_content_hashes"] = { str(Path(path).name): _sha256(Path(path)) for path in updated_shards } index["object_embedding_content_hash"] = _sha256(embed_path) index["content_hash"] = _content_hash(index) index_path.write_text(json.dumps(index, indent=2, sort_keys=True) + "\n") return { "split": split, "index": str(index_path), "rows": int(counters["rows"]), "rows_with_embedding": int(counters["rows_with_embedding"]), "missing_observation_ref": int(counters["missing_observation_ref"]), "unique_observation_refs": int(embedding_matrix.shape[0]), "embedding_path": str(embed_path), "embedding_content_hash": index["object_embedding_content_hash"], "content_hash_after": index["content_hash"], "split_hash": index.get("split_hash"), "task_counts": dict(sorted(task_counts.items())), } def _embedding_for_ref(key: tuple[str, str], h5_cache: dict[Path, Any]) -> np.ndarray: source_dataset, observation_ref = key archive_name, dataset_name, row_index = _parse_observation_ref(observation_ref) archive_path = Path(source_dataset) / archive_name if archive_path not in h5_cache: try: import h5py except ImportError as exc: # pragma: no cover raise ImportError("export_chart_object_embeddings.py requires h5py") from exc h5_cache[archive_path] = h5py.File(archive_path, "r") payload = np.asarray(h5_cache[archive_path][dataset_name][row_index], dtype=np.uint8) return _object_layout_embedding(payload.tobytes()) def _object_layout_embedding(jpeg_bytes: bytes) -> np.ndarray: try: from PIL import Image except ImportError as exc: # pragma: no cover raise ImportError("export_chart_object_embeddings.py requires Pillow") from exc image = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB").resize((96, 96)) arr = np.asarray(image, dtype=np.float32) / 255.0 gray = arr.mean(axis=2) saturation = arr.max(axis=2) - arr.min(axis=2) gy = np.zeros_like(gray) gx = np.zeros_like(gray) gy[1:, :] = np.abs(gray[1:, :] - gray[:-1, :]) gx[:, 1:] = np.abs(gray[:, 1:] - gray[:, :-1]) edge = gx + gy score = saturation + 0.5 * edge + 0.5 * np.abs(gray - float(np.median(gray))) threshold = max(float(np.quantile(score, 0.75)), float(score.mean() + 0.25 * score.std())) mask = score >= threshold components = _connected_components(mask) components = [component for component in components if len(component[0]) >= 8] components.sort(key=lambda component: len(component[0]), reverse=True) features: list[float] = [] for ys, xs in components[:4]: features.extend(_component_features(arr, gray, saturation, edge, ys, xs)) while len(features) < OBJECT_LAYOUT_EMBED_DIM: features.append(0.0) output = np.asarray(features[:OBJECT_LAYOUT_EMBED_DIM], dtype=np.float32) if output.shape[0] != OBJECT_LAYOUT_EMBED_DIM: raise AssertionError(f"expected {OBJECT_LAYOUT_EMBED_DIM} dims, got {output.shape[0]}") return output def _connected_components(mask: np.ndarray) -> list[tuple[np.ndarray, np.ndarray]]: height, width = mask.shape visited = np.zeros_like(mask, dtype=bool) components: list[tuple[np.ndarray, np.ndarray]] = [] for start_y in range(height): for start_x in range(width): if not mask[start_y, start_x] or visited[start_y, start_x]: continue stack = [(start_y, start_x)] visited[start_y, start_x] = True ys: list[int] = [] xs: list[int] = [] while stack: y, x = stack.pop() ys.append(y) xs.append(x) for next_y, next_x in ( (y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1), ): if ( 0 <= next_y < height and 0 <= next_x < width and mask[next_y, next_x] and not visited[next_y, next_x] ): visited[next_y, next_x] = True stack.append((next_y, next_x)) components.append((np.asarray(ys, dtype=np.int32), np.asarray(xs, dtype=np.int32))) return components def _component_features( arr: np.ndarray, gray: np.ndarray, saturation: np.ndarray, edge: np.ndarray, ys: np.ndarray, xs: np.ndarray, ) -> list[float]: height, width, _ = arr.shape pixels = arr[ys, xs] y_norm = ys.astype(np.float32) / max(float(height - 1), 1.0) x_norm = xs.astype(np.float32) / max(float(width - 1), 1.0) bbox_w = (float(xs.max() - xs.min() + 1) / float(width)) if xs.size else 0.0 bbox_h = (float(ys.max() - ys.min() + 1) / float(height)) if ys.size else 0.0 return [ float(xs.size) / float(height * width), float(2.0 * x_norm.mean() - 1.0), float(2.0 * y_norm.mean() - 1.0), float(2.0 * x_norm.std()), float(2.0 * y_norm.std()), bbox_w, bbox_h, *pixels.mean(axis=0).astype(float).tolist(), *pixels.std(axis=0).astype(float).tolist(), float(gray[ys, xs].mean()), float(saturation[ys, xs].mean()), float(edge[ys, xs].mean()), ] def _parse_observation_ref(value: str) -> tuple[str, str, int]: if "#" not in value: raise ValueError(f"invalid observation_ref: {value}") archive_name, ref = value.split("#", 1) parts = [part for part in ref.split("/") if part] if len(parts) != 2: raise ValueError(f"invalid observation_ref: {value}") return archive_name, parts[0], int(parts[1]) def _json_loads(value: str) -> dict[str, Any]: try: payload = json.loads(value) except json.JSONDecodeError: return {} return payload if isinstance(payload, dict) else {} def _metrics_by_task(rows: list[dict[str, Any]]) -> str: payload: dict[str, dict[str, int]] = {} for row in rows: for task, count in row["task_counts"].items(): payload.setdefault(task, {})[row["split"]] = int(count) return json.dumps(payload, indent=2, sort_keys=True) def _table(rows: list[dict[str, Any]]) -> str: lines = [ "% Auto-generated by scripts/export_chart_object_embeddings.py", "\\begin{tabular}{lrrr}", "\\toprule", "Split & Rows & With object embed & Unique refs \\\\", "\\midrule", ] for row in rows: lines.append( f"{row['split']} & {row['rows']} & " f"{row['rows_with_embedding']} & {row['unique_observation_refs']} \\\\" ) lines.extend(["\\bottomrule", "\\end{tabular}"]) return "\n".join(lines) def _report(metrics: dict[str, Any]) -> str: lines = [ "# Chart Object-Layout Embedding Export", "", "Decoded deployment-visible RGB observation refs into deterministic 64D " "foreground component/layout embeddings. No outcome, label, or hidden-branch " "fields are read.", "", "| Split | Rows | With embedding | Missing refs | Unique refs |", "| --- | ---: | ---: | ---: | ---: |", ] for row in metrics["splits"]: lines.append( f"| {row['split']} | {row['rows']} | {row['rows_with_embedding']} | " f"{row['missing_observation_ref']} | {row['unique_observation_refs']} |" ) return "\n".join(lines) def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None: (out_dir / "config.yaml").write_text( "\n".join(f"{key}: {value}" for key, value in sorted(vars(args).items())) + "\n" ) (out_dir / "command.txt").write_text( "python scripts/export_chart_object_embeddings.py " + " ".join(sys.argv[1:]) + "\n" ) (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n") hashes = {} for index_path in args.indexes: if index_path.exists(): index = json.loads(index_path.read_text()) hashes[str(index_path)] = { "content_hash": index.get("content_hash"), "split_hash": index.get("split_hash"), } (out_dir / "data_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n") (out_dir / "split_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n") def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _content_hash(index: dict[str, Any]) -> str: payload = dict(index) payload.pop("content_hash", None) return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() def _run(command: list[str]) -> str: try: return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip() except (subprocess.CalledProcessError, FileNotFoundError): return "" if __name__ == "__main__": raise SystemExit(main())