scripts/generate_questions.py

"""Generate an Orienter question manifest from a screenshot directory."""

import argparse
import json
import os
import sys
import tempfile
from pathlib import Path


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

from approach.pipeline_utils import parse_orienter_image_name


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
DEFAULT_PROMPT = "Identify interactable elements."


def iter_images(images_dir: Path, recursive: bool = False):
    if not images_dir.is_dir():
        raise FileNotFoundError(f"Screenshot directory does not exist: {images_dir}")
    pattern = "**/*" if recursive else "*"
    keyed_paths = []
    seen_image_ids = {}
    for path in images_dir.glob(pattern):
        if not path.is_file() or path.suffix.lower() not in IMAGE_SUFFIXES:
            continue
        _, _, image_id = parse_orienter_image_name(path.name)
        relative = path.relative_to(images_dir).as_posix()
        if image_id in seen_image_ids:
            raise ValueError(
                f"Duplicate image_id {image_id} from {seen_image_ids[image_id]!r} and {relative!r}"
            )
        seen_image_ids[image_id] = relative
        keyed_paths.append((image_id, path))
    for _, path in sorted(keyed_paths, key=lambda item: (item[0], item[1].as_posix())):
        yield path


def build_questions(images_dir: Path, prompt: str, recursive: bool = False):
    questions = []
    for index, image_path in enumerate(iter_images(images_dir, recursive=recursive)):
        relative_image = image_path.relative_to(images_dir).as_posix()
        _, _, image_id = parse_orienter_image_name(image_path.name)
        questions.append(
            {
                "question_id": index,
                "image": relative_image,
                "image_id": image_id,
                "text": prompt,
            }
        )
    return questions


def _atomic_text_writer(path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    return tempfile.mkstemp(
        prefix=".orienter-",
        suffix=f"{path.suffix}.tmp",
        dir=path.parent,
    )


def write_jsonl(path: Path, records):
    descriptor, temporary = _atomic_text_writer(path)
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as file:
            for record in records:
                file.write(json.dumps(record, ensure_ascii=False) + "\n")
        os.replace(temporary, path)
    except Exception:
        if os.path.exists(temporary):
            os.unlink(temporary)
        raise


def write_metadata_template(path: Path, questions):
    app_ids = sorted(
        {
            parse_orienter_image_name(Path(question["image"]).name)[0]
            for question in questions
        },
        key=int,
    )
    payload = {
        app_id: {
            "app_name": "",
            "app_description": "",
        }
        for app_id in app_ids
    }
    descriptor, temporary = _atomic_text_writer(path)
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as file:
            json.dump(payload, file, indent=2, ensure_ascii=False)
            file.write("\n")
        os.replace(temporary, path)
    except Exception:
        if os.path.exists(temporary):
            os.unlink(temporary)
        raise


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--images-dir", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--recursive", action="store_true")
    parser.add_argument(
        "--metadata-template",
        type=Path,
        help="optional JSON skeleton for the app metadata cache consumed by run_vlm",
    )
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    questions = build_questions(args.images_dir, args.prompt, recursive=args.recursive)
    if not questions:
        raise SystemExit(f"No supported images found in {args.images_dir}")
    write_jsonl(args.output, questions)
    if args.metadata_template:
        write_metadata_template(args.metadata_template, questions)
    print(
        json.dumps(
            {
                "questions": len(questions),
                "output": str(args.output),
                "metadata_template": str(args.metadata_template) if args.metadata_template else None,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()