scripts/meal_macro_pipeline.py

"""In-process meal macro pipeline for CLI and UI (no subprocess)."""

from __future__ import annotations

import csv
import json
import os
import sys
import tempfile
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import yaml

SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

os.environ.setdefault("YOLO_CONFIG_DIR", str(PROJECT_ROOT / ".ultralytics"))
os.environ.setdefault("TORCH_HOME", str(PROJECT_ROOT / ".torch"))

from ultralytics import YOLO  # pyright: ignore[reportPrivateImportUsage]

from build_usda_macros import read_queries
from estimate_macros_from_segments import estimate_grams, estimate_macros, load_macro_table
from gemini_analysis_to_usda_queries import find_component_for_segment_class, load_fallback_queries
from gemini_meal_analyzer import call_gemini, extract_json_text
from segmentation_utils import load_data_yaml_class_names
from usda_fdc_client import FdcClient, extract_macros

DEFAULT_WEIGHTS = (
    PROJECT_ROOT / "runs" / "foodseg103_target" / "yolov8s_target_e10_w02" / "weights" / "best.pt"
)
DEFAULT_DATA = PROJECT_ROOT / "data" / "processed" / "foodseg103_target_yolo" / "data.yaml"
DEFAULT_PROJECT_QUERIES = PROJECT_ROOT / "data" / "nutrition" / "project_macro_queries.csv"
DEFAULT_SAMPLE_MACROS = PROJECT_ROOT / "data" / "nutrition" / "sample_macros_per_100g.csv"
DEFAULT_PORTIONS = PROJECT_ROOT / "configs" / "portion_assumptions.yaml"
DEFAULT_CACHE = PROJECT_ROOT / "data" / "nutrition" / "fdc_cache.json"

CLASS_COLORS_BGR = {
    "meat": (80, 80, 255),
    "rice": (80, 220, 255),
    "vegetables": (80, 220, 80),
}


def _to_numpy(tensor_or_array: Any) -> np.ndarray:
    if hasattr(tensor_or_array, "cpu"):
        return tensor_or_array.cpu().numpy()
    return np.asarray(tensor_or_array)


def load_yolo_model(weights: Path = DEFAULT_WEIGHTS) -> YOLO | None:
    """Load the segmentation model, or None if weights are missing."""
    if not weights.exists():
        return None
    return YOLO(str(weights))


def _yolo_model(weights: Path = DEFAULT_WEIGHTS) -> YOLO:
    model = load_yolo_model(weights)
    if model is None:
        raise FileNotFoundError(f"YOLO weights not found: {weights}")
    return model


def predict_segments(
    image_path: Path,
    *,
    weights: Path = DEFAULT_WEIGHTS,
    data_yaml: Path = DEFAULT_DATA,
    imgsz: int = 512,
    conf: float = 0.05,
    model: YOLO | None = None,
) -> tuple[dict[str, Any], np.ndarray]:
    """Run YOLO segmentation; return JSON payload and BGR overlay image."""
    class_names = load_data_yaml_class_names(data_yaml)
    image_bgr = cv2.imread(str(image_path))
    if image_bgr is None:
        raise RuntimeError(f"Could not read image: {image_path}")

    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    height, width = image_rgb.shape[:2]
    yolo = model or _yolo_model(weights)
    result = yolo.predict(image_rgb, imgsz=imgsz, conf=conf, verbose=False)[0]

    area_by_class = {name: 0 for name in class_names.values()}
    occupied = np.zeros((height, width), dtype=bool)
    overlay = image_bgr.copy()

    masks_obj = result.masks
    boxes_obj = result.boxes
    if masks_obj is not None and boxes_obj is not None:
        masks = _to_numpy(masks_obj.data)
        classes = _to_numpy(boxes_obj.cls).astype(int)
        confs = _to_numpy(boxes_obj.conf)

        for idx in np.argsort(confs):
            item = masks[idx]
            if item.shape != (height, width):
                item = cv2.resize(item, (width, height), interpolation=cv2.INTER_NEAREST)
            binary = item > 0.5
            class_name = class_names[int(classes[idx])]
            new_pixels = np.logical_and(binary, ~occupied)
            area_by_class[class_name] += int(new_pixels.sum())
            occupied |= binary

            if not new_pixels.any():
                continue
            color = CLASS_COLORS_BGR.get(class_name, (200, 200, 200))
            overlay[new_pixels] = (
                0.45 * overlay[new_pixels].astype(np.float32)
                + 0.55 * np.array(color, dtype=np.float32)
            ).astype(np.uint8)

    total_food_area = sum(area_by_class.values())
    segments = []
    for class_name, pixel_area in area_by_class.items():
        if pixel_area <= 0 or total_food_area <= 0:
            continue
        segments.append(
            {
                "class_name": class_name,
                "pixel_area": pixel_area,
                "area_fraction": round(pixel_area / total_food_area, 6),
            }
        )

    payload = {
        "image": str(image_path),
        "weights": str(weights),
        "conf": conf,
        "total_food_area": total_food_area,
        "segments": segments,
    }
    return payload, overlay


def _analysis_to_query_rows(
    analysis: dict[str, Any],
    segment_classes: list[str],
    fallback_queries: dict[str, str],
) -> list[dict[str, str]]:
    best_by_class: dict[str, dict[str, Any]] = {}
    components = analysis.get("components", [])

    for component in components:
        class_name = str(component["class_name"])
        fdc_query = component.get("fdc_query") or component.get("likely_food") or class_name
        confidence = float(component.get("confidence", 0.0))
        current = best_by_class.get(class_name)
        if current is None or confidence > float(current["confidence"]):
            best_by_class[class_name] = {
                "class_name": class_name,
                "fdc_query": str(fdc_query),
                "confidence": confidence,
            }

    for segment_class in segment_classes:
        if segment_class in best_by_class:
            continue
        corrected = find_component_for_segment_class(segment_class, components)
        if corrected is not None:
            best_by_class[segment_class] = {
                "class_name": segment_class,
                "fdc_query": str(
                    corrected.get("fdc_query")
                    or corrected.get("likely_food")
                    or segment_class
                ),
                "confidence": float(corrected.get("confidence", 0.0)) - 0.01,
            }
        else:
            best_by_class[segment_class] = {
                "class_name": segment_class,
                "fdc_query": fallback_queries.get(segment_class, segment_class),
                "confidence": 0.0,
            }

    return [
        {"class_name": str(row["class_name"]), "fdc_query": str(row["fdc_query"])}
        for row in best_by_class.values()
    ]


def build_macros_from_queries(
    query_rows: list[dict[str, str]],
    *,
    api_key: str,
    cache_path: Path = DEFAULT_CACHE,
) -> list[dict[str, Any]]:
    client = FdcClient(api_key=api_key, cache_path=cache_path)
    output_rows: list[dict[str, Any]] = []
    for row in query_rows:
        class_name = row["class_name"]
        query = row["fdc_query"]
        food = client.search_first(query)
        macros = extract_macros(food)
        output_rows.append(
            {
                "class_name": class_name,
                "fdc_query": query,
                "fdc_id": food.get("fdcId") if food else "",
                "fdc_description": food.get("description") if food else "",
                **macros,
            }
        )
    client.save_cache()
    return output_rows


def _write_macros_csv(rows: list[dict[str, Any]], path: Path) -> None:
    fieldnames = [
        "class_name",
        "fdc_query",
        "fdc_id",
        "fdc_description",
        "kcal",
        "protein",
        "fat",
        "carbs",
    ]
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def _load_portions(total_plate_grams: float | None) -> dict[str, Any]:
    portions = yaml.safe_load(DEFAULT_PORTIONS.read_text(encoding="utf-8"))
    if total_plate_grams is not None:
        portions = {**portions, "total_plate_grams": float(total_plate_grams)}
    return portions


def run_meal_analysis(
    image_path: Path,
    *,
    use_gemini: bool = True,
    use_usda: bool = True,
    total_plate_grams: float | None = None,
    gemini_api_key: str | None = None,
    fdc_api_key: str | None = None,
    yolo_model: YOLO | None = None,
    yolo_conf: float = 0.05,
    yolo_imgsz: int = 512,
) -> dict[str, Any]:
    """Full pipeline: YOLO -> optional Gemini -> USDA -> macro estimate."""
    image_path = image_path.resolve()
    gemini_api_key = gemini_api_key or os.getenv("GEMINI_API_KEY")
    fdc_api_key = fdc_api_key or os.getenv("FDC_API_KEY")

    segments_payload, overlay_bgr = predict_segments(
        image_path,
        model=yolo_model,
        conf=yolo_conf,
        imgsz=yolo_imgsz,
    )
    segments: list[dict[str, Any]] = segments_payload["segments"]

    gemini_analysis: dict[str, Any] | None = None
    gemini_error: str | None = None
    query_rows = read_queries(DEFAULT_PROJECT_QUERIES)
    macro_rows: list[dict[str, Any]] | None = None

    if use_gemini and gemini_api_key:
        try:
            raw = call_gemini(gemini_api_key, "gemini-2.5-flash", image_path, segments_payload)
            gemini_analysis = extract_json_text(raw)
            segment_classes = [str(s["class_name"]) for s in segments]
            fallback = load_fallback_queries(DEFAULT_PROJECT_QUERIES)
            query_rows = _analysis_to_query_rows(gemini_analysis, segment_classes, fallback)
        except Exception as exc:
            gemini_error = str(exc)
    elif use_gemini:
        gemini_error = "GEMINI_API_KEY is not set."

    if use_usda and fdc_api_key:
        with tempfile.TemporaryDirectory() as tmp:
            macros_path = Path(tmp) / "macros.csv"
            macro_rows = build_macros_from_queries(query_rows, api_key=fdc_api_key)
            _write_macros_csv(macro_rows, macros_path)
            macro_table = load_macro_table(macros_path)
    else:
        if use_usda:
            msg = "FDC_API_KEY is not set; using sample macros."
            gemini_error = f"{gemini_error} {msg}".strip() if gemini_error else msg
        macro_table = load_macro_table(DEFAULT_SAMPLE_MACROS)
        macro_rows = [
            {"class_name": name, **values} for name, values in macro_table.items()
        ]

    portions = _load_portions(total_plate_grams)
    grams_by_class = estimate_grams(segments, portions)
    items, totals = estimate_macros(grams_by_class, macro_table)

    if macro_rows is None:
        macro_rows = [{"class_name": name, **values} for name, values in macro_table.items()]

    return {
        "image": str(image_path),
        "segments": segments_payload,
        "overlay_bgr": overlay_bgr,
        "gemini_analysis": gemini_analysis,
        "gemini_error": gemini_error,
        "usda_queries": query_rows,
        "macro_estimate": {"items": items, "totals": totals},
        "macros_per_100g": macro_rows,
    }