scripts/export_chart_object_embeddings.py

#!/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())