Spaces:
Sleeping
Sleeping
File size: 5,459 Bytes
11ab88a | 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 166 167 168 169 170 171 172 | from io import BytesIO
from pathlib import Path
from time import perf_counter
from fastapi import HTTPException, UploadFile
BASE_DIR = Path(__file__).resolve().parents[3]
MODEL_DIR = BASE_DIR / "ai" / "model"
MODEL_PATH = MODEL_DIR / "keras_model.h5"
LABELS_PATH = MODEL_DIR / "labels.txt"
MAX_FILE_SIZE = 10 * 1024 * 1024
ALLOWED_CONTENT_TYPES = {"image/jpeg", "image/png"}
_model = None
_labels: list[str] | None = None
def _load_labels() -> list[str]:
if not LABELS_PATH.exists():
raise HTTPException(status_code=500, detail="labels.txt file was not found.")
labels: list[str] = []
for line in LABELS_PATH.read_text(encoding="utf-8").splitlines():
label = line.strip()
if not label:
continue
parts = label.split(maxsplit=1)
labels.append(parts[1] if len(parts) == 2 and parts[0].isdigit() else label)
if not labels:
raise HTTPException(status_code=500, detail="labels.txt does not contain classes.")
return labels
def _load_model():
global _model
if _model is not None:
return _model
if not MODEL_PATH.exists():
raise HTTPException(status_code=500, detail="keras_model.h5 file was not found.")
try:
from tf_keras.layers import DepthwiseConv2D
from tf_keras.models import load_model
except ImportError:
try:
from tensorflow.keras.layers import DepthwiseConv2D
from tensorflow.keras.models import load_model
except ImportError as exc:
raise HTTPException(
status_code=500,
detail="TensorFlow is not installed. Run pip install -r requirements.txt.",
) from exc
except Exception as exc:
raise HTTPException(
status_code=500,
detail=f"TensorFlow loader import failed: {exc}",
) from exc
class TeachableMachineDepthwiseConv2D(DepthwiseConv2D):
@classmethod
def from_config(cls, config):
config.pop("groups", None)
return super().from_config(config)
try:
_model = load_model(
MODEL_PATH,
compile=False,
custom_objects={"DepthwiseConv2D": TeachableMachineDepthwiseConv2D},
)
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Model load failed: {exc}") from exc
return _model
def _get_labels() -> list[str]:
global _labels
if _labels is None:
_labels = _load_labels()
return _labels
def get_labels() -> list[str]:
return _get_labels()
def _model_input_size(model) -> tuple[int, int]:
input_shape = model.input_shape
if isinstance(input_shape, list):
input_shape = input_shape[0]
height = input_shape[1] or 224
width = input_shape[2] or 224
return int(width), int(height)
def _preprocess_image(image_bytes: bytes, size: tuple[int, int]):
try:
import numpy as np
from PIL import Image, ImageOps
except ImportError as exc:
raise HTTPException(
status_code=500,
detail="Image inference packages are missing. Run pip install -r requirements.txt.",
) from exc
try:
image = Image.open(BytesIO(image_bytes)).convert("RGB")
except Exception as exc:
raise HTTPException(status_code=400, detail="Uploaded file is not a readable image.") from exc
image = ImageOps.fit(image, size, method=Image.Resampling.LANCZOS)
image_array = np.asarray(image, dtype=np.float32)
normalized_image_array = (image_array / 127.5) - 1
return np.expand_dims(normalized_image_array, axis=0)
async def predict_image(file: UploadFile) -> dict:
if file.content_type not in ALLOWED_CONTENT_TYPES:
raise HTTPException(status_code=400, detail="Only jpg, jpeg, and png images are allowed.")
image_bytes = await file.read()
if not image_bytes:
raise HTTPException(status_code=400, detail="Uploaded file is empty.")
if len(image_bytes) > MAX_FILE_SIZE:
raise HTTPException(status_code=413, detail="Image files must be 10MB or smaller.")
model = _load_model()
labels = _get_labels()
model_input = _preprocess_image(image_bytes, _model_input_size(model))
started_at = perf_counter()
try:
predictions = model.predict(model_input, verbose=0)[0]
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Model inference failed: {exc}") from exc
inference_time_ms = round((perf_counter() - started_at) * 1000)
scores = predictions.tolist()
ranked = sorted(enumerate(scores), key=lambda item: item[1], reverse=True)
best_index, best_score = ranked[0]
def label_at(index: int) -> str:
return labels[index] if index < len(labels) else f"class_{index}"
return {
"success": True,
"filename": file.filename,
"predicted_class": label_at(best_index),
"confidence": float(best_score),
"top_k": [
{"label": label_at(index), "score": float(score)}
for index, score in ranked[: min(5, len(ranked))]
],
"model": {
"name": MODEL_PATH.name,
"version": "v1",
},
"inference_time_ms": inference_time_ms,
}
|