scripts/run_batch_meal_macro_demo.py

"""Run the end-to-end macro demo for a small batch of meal images.

Example:
    python scripts/run_batch_meal_macro_demo.py --image-glob "data/processed/foodseg103_target_yolo/images/val/*.jpg" --limit 5
"""

from __future__ import annotations

import argparse
import csv
import glob
import json
import os
import subprocess
import sys
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "output" / "demo"
DEFAULT_SUMMARY_CSV = PROJECT_ROOT / "output" / "batch_summary.csv"


def flatten_summary(summary_path: Path) -> dict[str, object]:
    """Flatten one demo summary.json into a single CSV-friendly row."""
    summary = json.loads(Path(summary_path).read_text(encoding="utf-8"))
    row: dict[str, object] = {
        "image": summary.get("image"),
        "summary_path": str(summary_path),
    }

    segments = (summary.get("segments") or {}).get("segments") or []
    for segment in segments:
        name = segment.get("class_name")
        if name:
            row[f"{name}_area_fraction"] = segment.get("area_fraction")

    gemini = summary.get("gemini_analysis") or {}
    if gemini.get("meal_summary") is not None:
        row["meal_summary"] = gemini.get("meal_summary")
    for component in gemini.get("components") or []:
        name = component.get("class_name")
        if name:
            row[f"{name}_likely_food"] = component.get("likely_food")
            row[f"{name}_fdc_query"] = component.get("fdc_query")

    totals = (summary.get("macro_estimate") or {}).get("totals") or {}
    for key in ("grams", "kcal", "protein", "fat", "carbs"):
        if key in totals:
            row[key] = totals[key]

    return row


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image-glob", required=True)
    parser.add_argument("--limit", type=int, default=5)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
    parser.add_argument("--summary-csv", type=Path, default=DEFAULT_SUMMARY_CSV)
    parser.add_argument("--use-gemini", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--use-usda", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args()


def run_one(image_path: Path, args: argparse.Namespace, env: dict[str, str]) -> Path:
    command = [
        sys.executable,
        "scripts/run_meal_macro_demo.py",
        "--image",
        str(image_path),
        "--output-dir",
        str(args.output_dir),
    ]
    if not args.use_gemini:
        command.append("--no-use-gemini")
    if not args.use_usda:
        command.append("--no-use-usda")

    print(f"+ {' '.join(command)}", flush=True)
    subprocess.run(command, cwd=PROJECT_ROOT, env=env, check=True)
    return args.output_dir / image_path.stem / "summary.json"


def main() -> None:
    args = parse_args()
    image_paths = [Path(path) for path in sorted(glob.glob(args.image_glob))]
    if args.limit is not None:
        image_paths = image_paths[: args.limit]
    if not image_paths:
        raise SystemExit(f"No images matched: {args.image_glob}")

    env = os.environ.copy()
    rows = []
    for index, image_path in enumerate(image_paths, start=1):
        print(f"\n[{index}/{len(image_paths)}] {image_path}", flush=True)
        summary_path = run_one(image_path, args, env)
        rows.append(flatten_summary(summary_path))

    fieldnames = sorted({key for row in rows for key in row})
    args.summary_csv.parent.mkdir(parents=True, exist_ok=True)
    with args.summary_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"Wrote batch summary: {args.summary_csv}")


if __name__ == "__main__":
    main()