scripts/smoke_control_path.py

"""Exercise the public one-image control path without APIs or model weights."""

import json
import sys
import tempfile
from pathlib import Path

from PIL import Image


if __package__ in {None, ""}:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from approach.run_ape import build_parser as build_ape_parser
from approach.run_ape import run as run_ape
from approach.run_vlm import build_parser as build_vlm_parser
from approach.run_vlm import run as run_vlm
from scripts.generate_questions import build_questions, write_jsonl


def main():
    with tempfile.TemporaryDirectory(prefix="orienter-smoke-") as tmpdir:
        root = Path(tmpdir)
        images_dir = root / "images"
        images_dir.mkdir()
        Image.new("RGB", (32, 32), color=(32, 64, 96)).save(images_dir / "123_4.png")

        questions_path = root / "questions.jsonl"
        write_jsonl(questions_path, build_questions(images_dir, "Smoke test"))

        candidates_path = root / "candidates.jsonl"
        vlm_args = build_vlm_parser().parse_args(
            [
                "--questions",
                str(questions_path),
                "--images-dir",
                str(images_dir),
                "--output",
                str(candidates_path),
            ]
        )

        def offline_processor(profile, question, image_path, ablation, key_index):
            return {"objects": {"button": "synthetic blue square"}}

        vlm_report = run_vlm(vlm_args, processor=offline_processor)

        predictions_path = root / "predictions.json"
        ape_args = build_ape_parser().parse_args(
            [
                "--questions",
                str(questions_path),
                "--candidates",
                str(candidates_path),
                "--images-dir",
                str(images_dir),
                "--output",
                str(predictions_path),
            ]
        )

        def offline_inference(**kwargs):
            return [
                {
                    "category_name": "button",
                    "bbox": [4, 5, 12, 10],
                    "score": 0.9,
                }
            ]

        ape_report = run_ape(ape_args, inference=offline_inference)
        predictions = json.loads(predictions_path.read_text(encoding="utf-8"))
        expected = {
            "image_id": 123004,
            "category_id": "button",
            "category_name": "button",
            "bbox": [4, 5, 12, 10],
        }
        if len(predictions) != 1 or any(
            predictions[0].get(key) != value for key, value in expected.items()
        ):
            raise RuntimeError(f"Unexpected smoke prediction: {predictions}")

        print(
            json.dumps(
                {
                    "status": "ok",
                    "vlm_records": vlm_report["completed"],
                    "ape_predictions": ape_report["predictions"],
                    "image_id": predictions[0]["image_id"],
                },
                indent=2,
            )
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())