multi / app.py
rishab1090's picture
Update app.py
e7b9294 verified
Raw
History Blame
1.72 kB
from fastapi import FastAPI, UploadFile, File, HTTPException, Header
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
import tensorflow as tf
import io
import os
app = FastAPI()
# Load your model
import os
import keras
from keras.models import load_model
# Read the model file into memory
with open("2.keras", "rb") as f:
byte_data = f.read()
# Wrap in a BytesIO object
model_file = io.BytesIO(byte_data)
# Load the model, giving a valid .keras filename
model = tf.keras.models.load_model(("2.keras", model_file))
# Now load the model
CLASS_NAMES = ['Fungi', 'Healthy', 'Nematode', 'Pest', 'Phytopthora', 'Virus']
# Define your API key (keep it secret in prod)
API_KEY = "mysecretkey"
@app.post("/predict")
async def predict(file: UploadFile = File(...), x_api_key: str = Header(None)):
if x_api_key != API_KEY:
raise HTTPException(status_code=401, detail="Invalid or missing API Key")
try:
contents = await file.read()
# Process the image
image = Image.open(io.BytesIO(contents)).convert("RGB")
image = image.resize((224, 224))
img_array = np.array(image).astype("float32")
img_array = np.expand_dims(img_array, axis=0)
# Predict
prediction = model.predict(img_array)
predicted_class = int(np.argmax(prediction[0]))
predicted_label = CLASS_NAMES[predicted_class]
return {
"prediction": predicted_label,
"probabilities": {
CLASS_NAMES[i]: float(round(prediction[0][i], 4)) for i in range(6)
}
}
except Exception as e:
return JSONResponse(status_code=500, content={"error": str(e)})