|
|
from fastapi import FastAPI |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from huggingface_hub import hf_hub_download |
|
|
import tensorflow as tf |
|
|
import os |
|
|
from fastapi import File, UploadFile |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
|
|
|
repo_id = "Sathvika-Alla/masterclass-2025" |
|
|
hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir="./model") |
|
|
hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir="./model") |
|
|
hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir="./model") |
|
|
|
|
|
|
|
|
model = tf.keras.models.load_model("./model") |
|
|
app = FastAPI() |
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=["*"], |
|
|
allow_credentials=True, |
|
|
allow_methods=["*"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
|
|
|
ANIMALS = ['Cat', 'Dog', 'Panda'] |
|
|
|
|
|
|
|
|
@app.post('/upload/image') |
|
|
async def uploadImage(img: UploadFile = File(...)): |
|
|
original_image = Image.open(img.file) |
|
|
resized_image = original_image.resize((64, 64)) |
|
|
images_to_predict = np.expand_dims(np.array(resized_image), axis=0) |
|
|
predictions = model.predict(images_to_predict) |
|
|
prediction_probabilities = predictions |
|
|
classifications = prediction_probabilities.argmax(axis=1) |
|
|
|
|
|
return ANIMALS[classifications.tolist()[0]] |