refactor: replace Python ML service runtime with Rust
This commit is contained in:
Asep Haryana Saputra
@@ -1,4 +1,4 @@
|
||||
MODEL_PATH=../../Machine_Learning/best_model/best_model.keras
|
||||
MODEL_PATH=../../Machine_Learning/model/model.onnx
|
||||
MODEL_INPUT_SIZE=224
|
||||
ML_SERVICE_HOST=0.0.0.0
|
||||
ML_SERVICE_PORT=8001
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
from fastapi import FastAPI, File, HTTPException, UploadFile
|
||||
|
||||
from model import ImageDecodeError, LABELS, SERVICE_NAME, SERVICE_VERSION, model_service
|
||||
from schemas import HealthResponse, MetadataResponse, PredictionResponse
|
||||
|
||||
|
||||
app = FastAPI(title="ZeaVis ML Service", version=SERVICE_VERSION)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def load_model() -> None:
|
||||
model_service.load()
|
||||
|
||||
|
||||
@app.get("/health", response_model=HealthResponse)
|
||||
def health() -> HealthResponse:
|
||||
return HealthResponse(status="ok", model_loaded=model_service.model_loaded)
|
||||
|
||||
|
||||
@app.get("/metadata", response_model=MetadataResponse)
|
||||
def metadata() -> MetadataResponse:
|
||||
return MetadataResponse(
|
||||
service_name=SERVICE_NAME,
|
||||
service_version=SERVICE_VERSION,
|
||||
model_path=str(model_service.model_path),
|
||||
model_loaded=model_service.model_loaded,
|
||||
input_size=model_service.input_size,
|
||||
labels=LABELS,
|
||||
)
|
||||
|
||||
|
||||
@app.post("/predict", response_model=PredictionResponse)
|
||||
async def predict(file: UploadFile = File(...)) -> PredictionResponse:
|
||||
if file.content_type is None or not file.content_type.startswith("image/"):
|
||||
raise HTTPException(status_code=400, detail="Uploaded file must be an image")
|
||||
|
||||
if not model_service.model_loaded:
|
||||
raise HTTPException(status_code=503, detail="Model is not loaded")
|
||||
|
||||
image_bytes = await file.read()
|
||||
|
||||
try:
|
||||
label, confidence, probabilities = model_service.predict(image_bytes)
|
||||
except ImageDecodeError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except Exception as exc:
|
||||
import logging
|
||||
|
||||
logging.exception("Prediction failed")
|
||||
raise HTTPException(status_code=500, detail="Prediction failed") from exc
|
||||
|
||||
return PredictionResponse(
|
||||
label=label,
|
||||
confidence=confidence,
|
||||
probabilities=probabilities,
|
||||
)
|
||||
@@ -1,72 +0,0 @@
|
||||
from io import BytesIO
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
import tensorflow as tf
|
||||
|
||||
|
||||
LABELS = ["Bercak Daun", "Daun Sehat", "Karat Daun", "Hawar Daun"]
|
||||
SERVICE_NAME = "zeavis-ml-service"
|
||||
SERVICE_VERSION = "0.1.0"
|
||||
|
||||
|
||||
class ImageDecodeError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class ModelService:
|
||||
def __init__(self) -> None:
|
||||
self.input_size = int(os.getenv("MODEL_INPUT_SIZE", "224"))
|
||||
self.model_path = self._resolve_model_path(os.getenv("MODEL_PATH", "../../Machine_Learning/best_model/best_model.keras"))
|
||||
self.model: tf.keras.Model | None = None
|
||||
self.load_error: str | None = None
|
||||
|
||||
def _resolve_model_path(self, model_path: str) -> Path:
|
||||
path = Path(model_path)
|
||||
if path.is_absolute():
|
||||
return path
|
||||
return (Path(__file__).resolve().parent / path).resolve()
|
||||
|
||||
@property
|
||||
def model_loaded(self) -> bool:
|
||||
return self.model is not None
|
||||
|
||||
def load(self) -> None:
|
||||
try:
|
||||
self.model = tf.keras.models.load_model(self.model_path, compile=False)
|
||||
self.load_error = None
|
||||
except Exception as exc:
|
||||
self.model = None
|
||||
self.load_error = str(exc)
|
||||
logging.exception("Failed to load ML model from %s", self.model_path)
|
||||
|
||||
def preprocess(self, image_bytes: bytes) -> np.ndarray:
|
||||
try:
|
||||
image = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
except (UnidentifiedImageError, OSError) as exc:
|
||||
raise ImageDecodeError("Uploaded file is not a valid image") from exc
|
||||
|
||||
image = image.resize((self.input_size, self.input_size))
|
||||
image_array = np.asarray(image, dtype=np.float32)
|
||||
return np.expand_dims(image_array, axis=0)
|
||||
|
||||
def predict(self, image_bytes: bytes) -> tuple[str, float, dict[str, float]]:
|
||||
if self.model is None:
|
||||
raise RuntimeError("Model is not loaded")
|
||||
|
||||
batch = self.preprocess(image_bytes)
|
||||
raw_predictions = self.model.predict(batch, verbose=0)[0]
|
||||
probabilities_array = np.asarray(raw_predictions, dtype=np.float32)
|
||||
top_index = int(np.argmax(probabilities_array))
|
||||
probabilities = {
|
||||
label: float(probabilities_array[index])
|
||||
for index, label in enumerate(LABELS)
|
||||
}
|
||||
|
||||
return LABELS[top_index], float(probabilities_array[top_index]), probabilities
|
||||
|
||||
|
||||
model_service = ModelService()
|
||||
@@ -1,10 +1,18 @@
|
||||
tasks:
|
||||
dev:
|
||||
command: .venv/bin/uvicorn main:app --host 0.0.0.0 --port 8001
|
||||
command: cargo run
|
||||
typecheck:
|
||||
command: python -m py_compile main.py model.py schemas.py
|
||||
command: cargo check
|
||||
inputs:
|
||||
- main.py
|
||||
- model.py
|
||||
- schemas.py
|
||||
- requirements.txt
|
||||
- Cargo.toml
|
||||
- src/**/*.rs
|
||||
test:
|
||||
command: cargo test
|
||||
inputs:
|
||||
- Cargo.toml
|
||||
- src/**/*.rs
|
||||
build:
|
||||
command: cargo build --release
|
||||
inputs:
|
||||
- Cargo.toml
|
||||
- src/**/*.rs
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
fastapi>=0.115.0
|
||||
uvicorn[standard]>=0.32.0
|
||||
tensorflow>=2.13.0
|
||||
pillow>=10.0.0
|
||||
python-multipart>=0.0.9
|
||||
@@ -1,21 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class HealthResponse(BaseModel):
|
||||
status: str
|
||||
model_loaded: bool
|
||||
|
||||
|
||||
class MetadataResponse(BaseModel):
|
||||
service_name: str
|
||||
service_version: str
|
||||
model_path: str
|
||||
model_loaded: bool
|
||||
input_size: int
|
||||
labels: list[str]
|
||||
|
||||
|
||||
class PredictionResponse(BaseModel):
|
||||
label: str
|
||||
confidence: float
|
||||
probabilities: dict[str, float]
|
||||
@@ -1,15 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from model import LABELS
|
||||
|
||||
|
||||
class ModelServiceTests(unittest.TestCase):
|
||||
def test_labels_match_training_class_order_with_display_names(self) -> None:
|
||||
self.assertEqual(
|
||||
LABELS,
|
||||
["Bercak Daun", "Daun Sehat", "Karat Daun", "Hawar Daun"],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user