File size: 6,212 Bytes
50a4916 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """Custom handler for Hugging Face Inference Endpoints.
Loads an EfficientNet-B0 Keras classifier (`classification_model.h5`) for the
Food-101-style food categories defined by `nutritional_database.json` and
returns the predicted dish along with calories / protein / fat / carbs.
"""
from __future__ import annotations
import base64
import io
import json
import logging
import os
from typing import Any, Dict, List, Optional, Union
from urllib.request import urlopen
import numpy as np
from PIL import Image
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
import tensorflow as tf # noqa: E402
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
IMAGE_SIZE = (224, 224)
TOP_K = 5
def _to_pil(payload: Any) -> Image.Image:
"""Best-effort conversion of whatever the endpoint received into a PIL image."""
if isinstance(payload, Image.Image):
return payload.convert("RGB")
if isinstance(payload, dict):
for key in ("image", "inputs", "url", "data", "bytes", "b64"):
if key in payload:
return _to_pil(payload[key])
raise ValueError(f"No image-like key found in dict input: {list(payload)}")
if isinstance(payload, (bytes, bytearray)):
return Image.open(io.BytesIO(bytes(payload))).convert("RGB")
if isinstance(payload, str):
# http(s) URL → fetch, otherwise assume base64 (with or without data URI prefix)
if payload.startswith(("http://", "https://")):
with urlopen(payload, timeout=15) as resp:
return Image.open(io.BytesIO(resp.read())).convert("RGB")
if payload.startswith("data:"):
payload = payload.split(",", 1)[-1]
try:
return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
except Exception as exc: # pragma: no cover - defensive
raise ValueError(f"String input is neither a URL nor valid base64: {exc}")
raise TypeError(f"Unsupported input type: {type(payload).__name__}")
def _preprocess(image: Image.Image) -> np.ndarray:
image = image.resize(IMAGE_SIZE, Image.BILINEAR)
arr = np.asarray(image, dtype=np.float32)
# EfficientNet's official preprocessor expects values in [0, 255]
arr = tf.keras.applications.efficientnet.preprocess_input(arr)
return np.expand_dims(arr, axis=0)
def _humanize(label: str) -> str:
return label.replace("_", " ").title()
class EndpointHandler:
"""Handler invoked by the HF Inference Endpoints toolkit on each request."""
def __init__(self, path: str = ""):
path = path or os.path.dirname(os.path.abspath(__file__))
logger.info("Loading food-recognition model from %s", path)
model_path = os.path.join(path, "classification_model.h5")
if not os.path.exists(model_path):
raise FileNotFoundError(f"classification_model.h5 not found at {model_path}")
self.model = tf.keras.models.load_model(model_path, compile=False)
nutri_path = os.path.join(path, "nutritional_database.json")
with open(nutri_path, "r", encoding="utf-8") as fh:
self.nutrition: Dict[str, Dict[str, float]] = json.load(fh)
# Class index → label. We assume the JSON ordering matches training-time
# label ordering (Food-101 is alphabetical, which the bundled DB respects).
self.labels: List[str] = list(self.nutrition.keys())
try:
output_shape = self.model.output_shape
num_outputs = output_shape[-1] if isinstance(output_shape, tuple) else None
if num_outputs is not None and num_outputs != len(self.labels):
logger.warning(
"Model output size (%s) != number of labels (%s); "
"predictions will be truncated/padded to the shorter length.",
num_outputs, len(self.labels),
)
except Exception: # pragma: no cover - some custom models lack output_shape
pass
logger.info("Loaded %d food classes", len(self.labels))
def _nutrition_for(self, label: str) -> Dict[str, Optional[float]]:
info = self.nutrition.get(label, {})
return {
"calories_per_100g": info.get("calories_per_100g"),
"protein_per_100g": info.get("protein_per_100g"),
"fat_per_100g": info.get("fat_per_100g"),
"carbs_per_100g": info.get("carbs_per_100g"),
"fiber_per_100g": info.get("fiber_per_100g"),
}
def __call__(self, data: Union[Dict[str, Any], bytes, Image.Image]) -> List[Dict[str, Any]]:
# The toolkit normally calls us with {"inputs": ..., "parameters": {...}};
# for raw image content-types it may pass bytes/PIL directly.
if isinstance(data, dict):
payload = data.get("inputs", data)
params = data.get("parameters", {}) or {}
else:
payload = data
params = {}
top_k = int(params.get("top_k", TOP_K))
image = _to_pil(payload)
batch = _preprocess(image)
preds = self.model.predict(batch, verbose=0)
if isinstance(preds, (list, tuple)):
preds = preds[0]
scores = np.asarray(preds).reshape(-1)
# Apply softmax only if the head clearly emits logits (not already in [0,1]).
if scores.min() < 0.0 or scores.max() > 1.0 + 1e-3:
scores = tf.nn.softmax(scores).numpy()
n = min(len(scores), len(self.labels))
scores = scores[:n]
labels = self.labels[:n]
order = np.argsort(scores)[::-1]
top_idx = order[: max(1, top_k)]
best_idx = int(top_idx[0])
best_label = labels[best_idx]
return [{
"dish": best_label,
"dish_name": _humanize(best_label),
"confidence": float(scores[best_idx]),
**self._nutrition_for(best_label),
"top_predictions": [
{
"dish": labels[int(i)],
"dish_name": _humanize(labels[int(i)]),
"confidence": float(scores[int(i)]),
}
for i in top_idx
],
}]
|