| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from PIL import Image |
| from transformers import SiglipVisionConfig, SiglipVisionModel |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| MODELS_DIR = ROOT / "models" |
|
|
|
|
| class NutritionHead(nn.Module): |
| def __init__(self, input_dim: int = 768) -> None: |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.LayerNorm(input_dim), |
| nn.Linear(input_dim, 768), |
| nn.GELU(), |
| nn.Dropout(0.15), |
| nn.Linear(768, 256), |
| nn.GELU(), |
| nn.Dropout(0.10), |
| nn.Linear(256, 5), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return F.softplus(self.net(x)) |
|
|
|
|
| @dataclass(frozen=True) |
| class CalorieKPrediction: |
| kcal: float |
| mass_g: float |
| protein_g: float |
| carbohydrate_g: float |
| fat_g: float |
| interval_90_radius_kcal: float |
| model: str |
| device: str |
|
|
|
|
| def prediction_to_dict(prediction: CalorieKPrediction) -> dict[str, Any]: |
| radius = prediction.interval_90_radius_kcal |
| return { |
| "kcal": round(prediction.kcal), |
| "mass_g": round(prediction.mass_g), |
| "protein_g": round(prediction.protein_g, 1), |
| "carbs_g": round(prediction.carbohydrate_g, 1), |
| "fat_g": round(prediction.fat_g, 1), |
| "interval_90_kcal": round(radius), |
| "kcal_low": max(0, round(prediction.kcal - radius)), |
| "kcal_high": round(prediction.kcal + radius), |
| "model": prediction.model, |
| "device": prediction.device, |
| "notes": "", |
| } |
|
|
|
|
| def _load_json(path: Path) -> dict: |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def preprocess_image(image: Image.Image, image_size: int = 224) -> torch.Tensor: |
| rgb = image.convert("RGB").resize((image_size, image_size), Image.Resampling.BICUBIC) |
| array = np.asarray(rgb, dtype=np.float32) / 255.0 |
| pixels = torch.from_numpy(array).permute(2, 0, 1) |
| return ((pixels - 0.5) / 0.5).unsqueeze(0) |
|
|
|
|
| def ensure_weights() -> Path: |
| checkpoint_path = MODELS_DIR / "config.json" |
| config = _load_json(checkpoint_path) |
| weights = MODELS_DIR / config["checkpoint_file"] |
| if weights.exists(): |
| return weights |
| from huggingface_hub import hf_hub_download |
|
|
| MODELS_DIR.mkdir(parents=True, exist_ok=True) |
| hf_hub_download( |
| repo_id="Dralkh/CalorieK", |
| filename=config["checkpoint_file"], |
| local_dir=str(MODELS_DIR), |
| ) |
| return weights |
|
|
|
|
| class CalorieK: |
| def __init__(self) -> None: |
| self.config = _load_json(MODELS_DIR / "config.json") |
| calibration = _load_json(MODELS_DIR / self.config["calibration_file"]) |
| self.interval = float( |
| self.config.get("interval_90_radius_kcal") or calibration.get("radius_kcal") or 118.95 |
| ) |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.image_size = int(self.config.get("image_size", 224)) |
|
|
| checkpoint_path = ensure_weights() |
|
|
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| vision_cfg = SiglipVisionConfig.from_json_file(str(MODELS_DIR / "vision_config.json")) |
| self.vision = SiglipVisionModel(vision_cfg) |
| vision_target = self.vision.vision_model if hasattr(self.vision, "vision_model") else self.vision |
| vision_target.load_state_dict(checkpoint["vision_state"]) |
|
|
| self.head = NutritionHead(int(self.config.get("hidden_size", 768))) |
| self.head.load_state_dict(checkpoint["head_state"]) |
| self.vision.to(self.device).eval() |
| self.head.to(self.device).eval() |
|
|
| @torch.inference_mode() |
| def predict(self, image: Image.Image) -> CalorieKPrediction: |
| pixels = preprocess_image(image, self.image_size).to(self.device) |
| out = self.vision(pixel_values=pixels) |
| feats = out.pooler_output if out.pooler_output is not None else out.last_hidden_state[:, 0] |
| kcal, mass_g, protein_g, carb_g, fat_g = self.head(feats.float())[0].tolist() |
| return CalorieKPrediction( |
| kcal=float(kcal), |
| mass_g=float(mass_g), |
| protein_g=float(protein_g), |
| carbohydrate_g=float(carb_g), |
| fat_g=float(fat_g), |
| interval_90_radius_kcal=self.interval, |
| model="plato", |
| device=str(self.device), |
| ) |
|
|
|
|
| _model: CalorieK | None = None |
|
|
|
|
| def get_model() -> CalorieK: |
| global _model |
| if _model is None: |
| _model = CalorieK() |
| return _model |
|
|
