IDS75912/CTAIAnimalClassifierFastApi

1import uvicorn
2
3
4import fastapi
5from fastapi import FastAPI, Request
6from fastapi.middleware.cors import CORSMiddleware
7from fastapi import File, UploadFile
8import numpy as np
9from PIL import Image
10
11
12from typing import Any, Dict
13import os
14import pkgutil
15
16from huggingface_hub import hf_hub_download
17from huggingface_hub import hf_hub_url
18import requests
19import tempfile
20import shutil
21from typing import Any, Dict
22
23import tensorflow as tf
24import traceback
25import logging
26from tensorflow import keras
27
28app = FastAPI(title="1.3 - AI Model Deployment - HF Hub + FastAPI",)
29''' browser: http://localhost:8000/docs'''
30
31from fastapi.middleware.cors import CORSMiddleware
32app.add_middleware(
33    CORSMiddleware,
34    allow_origins=["*"],
35    allow_credentials=True,
36    allow_methods=["*"],
37    allow_headers=["*"],
38)
39
40
41ANIMALS = ['Cat', 'Dog', 'Panda'] # Animal names here, these represent the labels of the images that we trained our model on.
42
43# 1) download your SavedModel from the Hub into a writable directory (Spaces often
44# take HF_MODEL_DIR or default model_dir).
45repo_id = "IDS75912/masterclass-2025"
46local_model_dir = os.environ.get('HF_MODEL_DIR', './model_dir')
47
48# Ensure the directory exists and is writable. If creating fails, raise a clear error.
49try:
50    os.makedirs(local_model_dir, exist_ok=True)
51except Exception as e:
52    raise RuntimeError(f"Cannot create model directory '{local_model_dir}'. Ensure the process has write access or set HF_MODEL_DIR to a writable path., Error: {e}")
53
54# download files into local_model_dir and load model with resilient error handling
55model = None
56model_load_error = None
57#try:
58    # First try using a cache dir so downloads happen in a shared cache and final move
59    # into local_model_dir is less likely to require risky renames inside the repo.
60    # cache_dir = '/tmp/.cache/huggingface'
61    # os.makedirs(cache_dir, exist_ok=True)
62
63hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir=local_model_dir )
64hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir=local_model_dir)
65hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir=local_model_dir)
66
67    # 2) load it
68model = tf.keras.models.load_model(local_model_dir)
69logging.info(f"Model loaded successfully from {local_model_dir}")
70
71
72@app.post('/upload/image')
73async def uploadImage(img: UploadFile = File(...)):
74    if model is None:
75        # Model isn't available — return a helpful error to the caller instead of crashing.
76        return fastapi.Response(status_code=503, content=f"Model not loaded: {model_load_error}")
77
78    original_image = Image.open(img.file) # Read the bytes and process as an image
79    if original_image.mode == 'RGBA':
80        original_image = original_image.convert('RGB')
81    resized_image = original_image.resize((64, 64)) # Resize
82    images_to_predict = np.expand_dims(np.array(resized_image), axis=0) # Our AI Model wanted a list of images, but we only have one, so we expand it's dimension
83    predictions = model.predict(images_to_predict) # The result will be a list with predictions in the one-hot encoded format: [ [0 1 0] ]
84    prediction_probabilities = predictions
85    classifications = prediction_probabilities.argmax(axis=1) # We try to fetch the index of the highest value in this list [ [1] ]
86
87    return ANIMALS[classifications.tolist()[0]] # Fetch the first item in our classifications array, format it as a list first, result will be e.g.: "Dog"
88
89@app.get("/")
90def read_root() -> Dict[str, Any]:
91    """Root endpoint."""
92    return {"message": "Hello from FastAPI in 1.3 - AI Model Deployment - HF Hub + FastAPI"}
93
94
95@app.get("/version")
96def versions() -> Dict[str, Any]:
97    """Return key package versions and whether TensorFlow is available."""
98    return {
99        "fastapi": fastapi.__version__,
100
101    }
102
103
104@app.get("/predict")
105def predict_stub() -> Dict[str, Any]:
106
107    # This is a stub, so we're not doing a real prediction
108    if model is None:
109        return {"prediction": "model not loaded", "error": model_load_error}
110    return {"prediction": "stub, we're not doing a real prediction"}
111
112
113
114if __name__ == "__main__":
115	# Run with: conda run -n gradio uvicorn main:app --reload
116	import uvicorn
117
118	uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True) # ? 7860 instead of 8000