src/predict.py

import torch

from .nutrition import get_nutrition


def predict_top_k(image, processor, model, calorie_db: dict, top_k: int = 3) -> list[dict]:
    """Run image classification and return top-k predictions with nutrition metadata."""
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]

    top = probs.topk(top_k)
    results = []
    for idx, score in zip(top.indices.tolist(), top.values.tolist()):
        raw_label = model.config.id2label[idx]
        nutrition = get_nutrition(raw_label, calorie_db)
        results.append(
            {
                "label": raw_label.replace("_", " ").title(),
                "raw_label": raw_label,
                "confidence": round(score * 100, 1),
                "category": nutrition["category"],
                "calories_per_100g": nutrition["calories_per_100g"],
                "serving_g": nutrition["serving_g"],
                "calories_per_serving": nutrition["calories_per_serving"],
            }
        )

    return results