Added /health endpoint

model-lab/app/main.py

@@ -1,7 +1,7 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, status
 from PIL import Image
 import io
-from app.model import predict_emotion
+from app.model import predict_emotion, is_model_ready
 
 app = FastAPI(title="Face Sentiment API")
 
@@ -10,3 +10,10 @@ async def predict(file: UploadFile = File(...)):
     image = Image.open(io.BytesIO(await file.read()))
     emotion = predict_emotion(image)
     return {"emotion": emotion}
+
+@app.get("/health", status_code=status.HTTP_200_OK)
+def get_health():
+    if is_model_ready():
+        return {"status" : "ok"}
+    else:
+        return {"status" : "unavailable"}

model-lab/app/model.py

@@ -1,11 +1,23 @@
+from typing import Optional
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 from PIL import Image
 import torch
 
 MODEL_NAME = "trpakov/vit-face-expression"
+MODEL_READY = False
 
-extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
-model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+try:
+    extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
+    model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+except Exception:
+    extractor = None
+    model = None
+    MODEL_READY = False
+else:
+    MODEL_READY = True
+
+def is_model_ready() -> bool:
+    return MODEL_READY
 
 def predict_emotion(image: Image.Image):
     inputs = extractor(images=image, return_tensors="pt")