Update app.py

app.py

@@ -1,100 +1,103 @@
-"""
-FastAPI service exposing BinhQuocNguyen/food‑recognition‑model.
-"""
-
 # ------------------------------------------------------------
-# 1️⃣ Imports
+# FastAPI service exposing BinhQuocNguyen/food-recognition-model
 # ------------------------------------------------------------
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import base64
-import io
-
+import base64, io
 from PIL import Image
-from transformers import pipeline
+import torch
+import numpy as np
 
-# ------------------------------------------------------------
-# 2️⃣ Load the model once at startup (fast, no re‑load per request)
-# ------------------------------------------------------------
-classifier = pipeline(
-    "image-classification",
-    model="BinhQuocNguyen/food-recognition-model"
-)
+# Transformers imports
+from transformers import AutoModel, AutoImageProcessor
 
-# ------------------------------------------------------------
-# 3️⃣ Simple in‑memory “nutritional DB”
-#    (calories per 100 g, average portion weight in grams)
-# ------------------------------------------------------------
+# -------------------------------------------------------------------
+# 1️⃣ Load the model & processor (once, at import time)
+# -------------------------------------------------------------------
+MODEL_NAME = "BinhQuocNguyen/food-recognition-model"
+
+# AutoModel knows the custom architecture (food_recognition) because
+# the repository ships a proper `config.json`.
+model = AutoModel.from_pretrained(MODEL_NAME)
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+
+# Put the model on CPU – the Space has no GPU.
+device = torch.device("cpu")
+model.to(device)
+model.eval()
+
+# Mapping from class index → readable label (comes from the config)
+id2label = model.config.id2label  # dict[int, str]
+
+# -------------------------------------------------------------------
+# 2️⃣ Minimal nutrient lookup table (extend as you like)
+# -------------------------------------------------------------------
 nutrient_db = {
     "Apple":  {"calories_per_100g": 52,  "portion_g": 182},
     "Banana": {"calories_per_100g": 89,  "portion_g": 118},
     "Orange": {"calories_per_100g": 43,  "portion_g": 131},
     "Pizza":  {"calories_per_100g": 266, "portion_g": 200},
     "Bread":  {"calories_per_100g": 265, "portion_g": 30},
-    # ← add the rest of your favourite foods here
+    # Add the rest of the 101 categories if you need them
 }
 
-# ------------------------------------------------------------
-# 4️⃣ Pydantic model for the incoming JSON payload
-# ------------------------------------------------------------
+# -------------------------------------------------------------------
+# 3️⃣ Pydantic model for the incoming JSON payload
+# -------------------------------------------------------------------
 class ImageRequest(BaseModel):
-    """Base64‑encoded image sent by the client."""
-    image: str
+    image: str   # base64‑encoded JPEG/PNG
 
-
-# ------------------------------------------------------------
-# 5️⃣ FastAPI app & health‑check endpoint
-# ------------------------------------------------------------
 app = FastAPI()
 
 
+# ------------------------------------------------------------
+# Health‑check endpoint (optional)
+# ------------------------------------------------------------
 @app.get("/")
-def root():
-    """Simple health‑check."""
+def health():
     return {"message": "Food‑Recognition API is up"}
 
 
 # ------------------------------------------------------------
-# 6️⃣ Main inference endpoint
+# 4️⃣ Main inference endpoint
 # ------------------------------------------------------------
 @app.post("/analyze")
 def analyze(request: ImageRequest):
-    """
-    1️⃣ Decode the base64 image.
-    2️⃣ Run the classifier.
-    3️⃣ Look up (or fall back to) nutritional information.
-    4️⃣ Return a JSON response.
-    """
-    # ---- 1️⃣ decode the base64 image ---------------------------------
+    # ---- 4.1 decode the base64 image ---------------------------------
     try:
-        raw_bytes = base64.b64decode(request.image)
-        pil_img = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
-    except Exception as e:  # pragma: no cover
-        raise HTTPException(status_code=400, detail="Invalid base64 image") from e
-
-    # ---- 2️⃣ run the classifier --------------------------------------
-    results = classifier(pil_img)
-    if not results:
-        raise HTTPException(status_code=500, detail="Model returned no results")
-
-    top = results[0]
-    label = top.get("label")
-    confidence = top.get("score")
-
-    # ---- 3️⃣ look up nutrition data ----------------------------------
-    # Fallback values if the label isn’t in the DB
+        raw = base64.b64decode(request.image)
+        pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid base64 image")
+
+    # ---- 4.2 preprocess ------------------------------------------------
+    inputs = processor(images=pil_img, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    # ---- 4.3 forward pass ---------------------------------------------
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # The model returns logits (shape [1, num_classes])
+    logits = outputs.logits.squeeze(0)          # [num_classes]
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+
+    # ---- 4.4 get top‑1 prediction --------------------------------------
+    top_idx = int(probs.argmax().item())
+    confidence = float(probs[top_idx].item())
+    label = id2label.get(top_idx, "unknown")
+
+    # ---- 4.5 lookup nutrition -----------------------------------------
     nutrition = nutrient_db.get(label, {"calories_per_100g": 0, "portion_g": 100})
     calories_per_100g = nutrition["calories_per_100g"]
     portion_g = nutrition["portion_g"]
+    estimated_calories = calories_per_100g * (portion_g / 100.0)
 
-    # ---- 4️⃣ estimated calories for the default portion size ----------
-    est_calories = calories_per_100g * (portion_g / 100.0)
-
-    # ---- 5️⃣ build the JSON response --------------------------------
+    # ---- 4.6 build JSON response ---------------------------------------
     return {
         "label": label,
         "confidence": confidence,
         "estimated_portion_g": portion_g,
         "calories_per_100g": calories_per_100g,
-        "estimated_calories": round(est_calories, 2),
+        "estimated_calories": round(estimated_calories, 2)
     }