visualize_self_localization_samples.py

#!/usr/bin/env python3
"""Create a portable HTML sampler for whereami JSONL records."""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
from typing import Any

try:
    from .common import html_escape_text, image_to_data_uri
except ImportError:
    from common import html_escape_text, image_to_data_uri


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Render a random whereami sample viewer.")
    parser.add_argument("--input-jsonl", type=Path, required=True)
    parser.add_argument("--output-html", type=Path, required=True)
    parser.add_argument("--examples", type=int, default=12)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--max-image-width", type=int, default=360)
    return parser.parse_args()


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    with path.open() as handle:
        return [json.loads(line) for line in handle if line.strip()]


def image_block(path: str, index: int, max_width: int) -> str:
    uri = image_to_data_uri(Path(path), max_width=max_width)
    if uri is None:
        return f"<div class='missing'>Image {index + 1}: missing<br>{html_escape_text(path)}</div>"
    return f"<figure><img src='{uri}' loading='lazy'><figcaption>Image {index + 1}: {html_escape_text(Path(path).name)}</figcaption></figure>"


def render_record(record: dict[str, Any], max_width: int) -> str:
    metadata = record.get("metadata") or {}
    anchor = metadata.get("anchor") or {}
    cell = record.get("ground_truth_camera_cell")
    facing = record.get("ground_truth_facing_vector")
    images = "".join(image_block(path, index, max_width) for index, path in enumerate(record.get("image_paths") or []))
    if not images:
        images = "<div class='missing'>No exported image paths in this record.</div>"
    diagnostics = {
        "group": record.get("image_group_id"),
        "anchor image": record.get("anchor_image_index"),
        "anchor frame": record.get("anchor_frame_index"),
        "sharpness": anchor.get("frame_sharpness"),
        "max cell-axis offset (m)": (anchor.get("camera_cell_margin") or {}).get("cell_center_axis_offset_m"),
        "facing bin margin (deg)": (anchor.get("facing_bin_margin") or {}).get("facing_bin_margin_deg"),
        "group visible grid instances": metadata.get("visible_grid_instance_ids_used_for_crop"),
        "group visible grid categories": metadata.get("visible_grid_categories_used_for_crop"),
        "anchor visible grid instances": anchor.get("visible_grid_instance_ids"),
    }
    diagnostic_rows = "".join(
        f"<tr><th>{html_escape_text(key)}</th><td>{html_escape_text(value)}</td></tr>" for key, value in diagnostics.items()
    )
    return f"""<article>
<h2>{html_escape_text(record.get('qa_id'))}</h2>
<div class='images'>{images}</div>
<div class='columns'>
<section><h3>Ground truth</h3><pre>CAMERA_CELL: {html_escape_text(cell)}\nFACING_VECTOR: {html_escape_text(facing)}</pre><table>{diagnostic_rows}</table></section>
<section><h3>ASCII object grid</h3><pre>{html_escape_text(record.get('grid_text'))}</pre></section>
</div>
<details><summary>Prompt</summary><pre>{html_escape_text(record.get('prompt'))}</pre></details>
</article>"""


def main() -> None:
    args = parse_args()
    records = read_jsonl(args.input_jsonl)
    rng = random.Random(args.seed)
    if args.examples > 0 and len(records) > args.examples:
        records = rng.sample(records, args.examples)
    body = "\n".join(render_record(record, args.max_image_width) for record in records)
    page = f"""<!doctype html>
<html><head><meta charset='utf-8'><title>WhereAmI samples</title>
<style>
body {{ margin: 0; background: #f4f5f7; color: #20252b; font: 14px/1.45 Arial, sans-serif; }}
main {{ max-width: 1440px; margin: auto; padding: 24px; }}
article {{ background: white; border: 1px solid #d7dce2; border-radius: 6px; margin: 0 0 20px; padding: 18px; }}
h1 {{ margin: 0 0 20px; }} h2 {{ margin: 0 0 12px; font-size: 17px; }} h3 {{ margin: 0 0 8px; font-size: 14px; }}
.images {{ display: flex; flex-wrap: wrap; gap: 12px; margin-bottom: 16px; }} figure {{ margin: 0; width: 360px; }} img {{ width: 100%; border: 1px solid #d7dce2; display: block; }} figcaption {{ color: #5d6670; font-size: 12px; }}
.columns {{ display: grid; grid-template-columns: minmax(280px, 0.75fr) minmax(320px, 1.25fr); gap: 18px; }}
pre {{ overflow: auto; background: #f7f8fa; border: 1px solid #e2e6ea; padding: 10px; margin: 0; white-space: pre; }}
table {{ border-collapse: collapse; width: 100%; margin-top: 10px; }} th, td {{ border: 1px solid #e2e6ea; padding: 5px; text-align: left; vertical-align: top; }} th {{ width: 42%; background: #f7f8fa; }}
details {{ margin-top: 14px; }} details pre {{ margin-top: 8px; white-space: pre-wrap; }} .missing {{ color: #9c2b2e; }}
@media (max-width: 800px) {{ main {{ padding: 12px; }} .columns {{ grid-template-columns: 1fr; }} figure {{ width: min(100%, 360px); }} }}
</style></head><body><main><h1>WhereAmI self-localization samples ({len(records)})</h1>{body}</main></body></html>"""
    args.output_html.parent.mkdir(parents=True, exist_ok=True)
    args.output_html.write_text(page)
    print(f"Wrote {len(records)} samples to {args.output_html}")


if __name__ == "__main__":
    main()