File size: 4,438 Bytes
89af506 27384f5 89af506 059d1d9 89af506 059d1d9 89af506 27384f5 678abfb 89af506 521c944 36dc52d 89af506 6b9607e 521c944 6b9607e 521c944 059d1d9 6b9607e 059d1d9 6b9607e 89af506 6b9607e 89af506 059d1d9 6b9607e 89af506 059d1d9 89af506 678abfb 89af506 27384f5 89af506 27384f5 89af506 059d1d9 27384f5 89af506 05de3e3 |
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 |
import uvicorn
import fastapi
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi import File, UploadFile
import numpy as np
from PIL import Image
from typing import Any, Dict
import os
import pkgutil
from huggingface_hub import hf_hub_download
from huggingface_hub import hf_hub_url
import requests
import tempfile
import shutil
from typing import Any, Dict
import tensorflow as tf
import traceback
import logging
from tensorflow import keras
app = FastAPI(title="1.3 - AI Model Deployment - HF Hub + FastAPI",)
''' browser: http://localhost:8000/docs'''
from fastapi.middleware.cors import CORSMiddleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
ANIMALS = ['Cat', 'Dog', 'Panda'] # Animal names here, these represent the labels of the images that we trained our model on.
# 1) download your SavedModel from the Hub into a writable directory (Spaces often
# take HF_MODEL_DIR or default model_dir).
repo_id = "IDS75912/masterclass-2025"
local_model_dir = os.environ.get('HF_MODEL_DIR', './model_dir')
# Ensure the directory exists and is writable. If creating fails, raise a clear error.
try:
os.makedirs(local_model_dir, exist_ok=True)
except Exception as e:
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}")
# download files into local_model_dir and load model with resilient error handling
model = None
model_load_error = None
#try:
# First try using a cache dir so downloads happen in a shared cache and final move
# into local_model_dir is less likely to require risky renames inside the repo.
# cache_dir = '/tmp/.cache/huggingface'
# os.makedirs(cache_dir, exist_ok=True)
hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir=local_model_dir )
hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir=local_model_dir)
hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir=local_model_dir)
# 2) load it
model = tf.keras.models.load_model(local_model_dir)
logging.info(f"Model loaded successfully from {local_model_dir}")
@app.post('/upload/image')
async def uploadImage(img: UploadFile = File(...)):
if model is None:
# Model isn't available — return a helpful error to the caller instead of crashing.
return fastapi.Response(status_code=503, content=f"Model not loaded: {model_load_error}")
original_image = Image.open(img.file) # Read the bytes and process as an image
if original_image.mode == 'RGBA':
original_image = original_image.convert('RGB')
resized_image = original_image.resize((64, 64)) # Resize
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
predictions = model.predict(images_to_predict) # The result will be a list with predictions in the one-hot encoded format: [ [0 1 0] ]
prediction_probabilities = predictions
classifications = prediction_probabilities.argmax(axis=1) # We try to fetch the index of the highest value in this list [ [1] ]
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"
@app.get("/")
def read_root() -> Dict[str, Any]:
"""Root endpoint."""
return {"message": "Hello from FastAPI in 1.3 - AI Model Deployment - HF Hub + FastAPI"}
@app.get("/version")
def versions() -> Dict[str, Any]:
"""Return key package versions and whether TensorFlow is available."""
return {
"fastapi": fastapi.__version__,
}
@app.get("/predict")
def predict_stub() -> Dict[str, Any]:
# This is a stub, so we're not doing a real prediction
if model is None:
return {"prediction": "model not loaded", "error": model_load_error}
return {"prediction": "stub, we're not doing a real prediction"}
if __name__ == "__main__":
# Run with: conda run -n gradio uvicorn main:app --reload
import uvicorn
uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True) # ? 7860 instead of 8000 |