SDK-Streamlit / scripts /gemini_analysis_to_usda_queries.py
Gilgarmesh's picture
Upload 22 files
e0e2b27 verified
Raw
History Blame Contribute Delete
4.1 kB
"""Convert Gemini meal analysis JSON into a USDA query CSV."""
from __future__ import annotations
import argparse
import csv
import json
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_ANALYSIS = PROJECT_ROOT / "data" / "nutrition" / "gemini_meal_analysis.json"
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "nutrition" / "gemini_usda_queries.csv"
DEFAULT_FALLBACK = PROJECT_ROOT / "data" / "nutrition" / "project_macro_queries.csv"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--analysis", type=Path, default=DEFAULT_ANALYSIS)
parser.add_argument("--segments", type=Path)
parser.add_argument("--fallback-queries", type=Path, default=DEFAULT_FALLBACK)
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
return parser.parse_args()
def load_segment_classes(path: Path | None) -> list[str]:
if path is None or not path.exists():
return []
payload = json.loads(path.read_text(encoding="utf-8"))
segments = payload.get("segments", payload if isinstance(payload, list) else [])
return [segment["class_name"] for segment in segments if "class_name" in segment]
def load_fallback_queries(path: Path) -> dict[str, str]:
if not path.exists():
return {}
with path.open(newline="", encoding="utf-8") as f:
return {row["class_name"]: row["fdc_query"] for row in csv.DictReader(f)}
def main() -> None:
args = parse_args()
analysis = json.loads(args.analysis.read_text(encoding="utf-8"))
segment_classes = load_segment_classes(args.segments)
fallback_queries = load_fallback_queries(args.fallback_queries)
best_by_class = {}
components = analysis.get("components", [])
for component in components:
class_name = 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 > current["confidence"]:
best_by_class[class_name] = {
"class_name": class_name,
"fdc_query": fdc_query,
"confidence": confidence,
}
for segment_class in segment_classes:
if segment_class in best_by_class:
continue
corrected_component = find_component_for_segment_class(segment_class, components)
if corrected_component is not None:
best_by_class[segment_class] = {
"class_name": segment_class,
"fdc_query": corrected_component.get("fdc_query")
or corrected_component.get("likely_food")
or segment_class,
"confidence": float(corrected_component.get("confidence", 0.0)) - 0.01,
}
continue
best_by_class[segment_class] = {
"class_name": segment_class,
"fdc_query": fallback_queries.get(segment_class, segment_class),
"confidence": 0.0,
}
rows = [
{"class_name": row["class_name"], "fdc_query": row["fdc_query"]}
for row in best_by_class.values()
]
args.output.parent.mkdir(parents=True, exist_ok=True)
with args.output.open("w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=["class_name", "fdc_query"])
writer.writeheader()
writer.writerows(rows)
print(f"Wrote USDA query CSV: {args.output}")
def find_component_for_segment_class(segment_class: str, components: list[dict]) -> dict | None:
needle = segment_class.lower()
for component in sorted(
components,
key=lambda item: float(item.get("confidence", 0.0)),
reverse=True,
):
text = " ".join(
str(component.get(field, ""))
for field in ("notes", "likely_food", "fdc_query")
).lower()
if needle in text:
return component
return None
if __name__ == "__main__":
main()