File size: 1,770 Bytes
d042a88 506acc6 d042a88 506acc6 d042a88 506acc6 d042a88 |
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 |
from fastapi import FastAPI
app = FastAPI()
from fastapi.middleware.cors import CORSMiddleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
from tensorflow import keras
import tensorflow as tf
import os
from huggingface_hub import hf_hub_download
ANIMALS = ['Cat', 'Dog', 'Panda'] # Animal names here, these represent the labels of the images that we trained our model on.
# Download model from Hugging Face Hub (avoids storing large file in repo)
model_path = hf_hub_download(
repo_id="benitovvt/animal-classification-model", # Replace with your model repo
filename="model.keras",
cache_dir="./model_cache"
)
model = tf.keras.models.load_model(model_path)
print(f"✓ Model loaded from Hugging Face Hub: {model_path}")
from fastapi import File, UploadFile
import numpy as np
from PIL import Image
@app.post('/upload/image')
async def uploadImage(img: UploadFile = File(...)):
original_image = Image.open(img.file) # Read the bytes and process as an image
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" |