HashamUllah's picture
Update app.py
fceec31 verified
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import json
app = FastAPI()
# Load the TensorFlow model
model = tf.keras.models.load_model('./plant_disease_detection.h5')
# Load categories
with open('./categories.json') as f:
categories = json.load(f)
def preprocess_image(image_bytes):
# Convert the image to a NumPy array
image = Image.open(io.BytesIO(image_bytes))
image = image.resize((224, 224)) # Adjust size as needed
image_array = np.array(image) / 255.0 # Normalize to [0, 1]
image_array = np.expand_dims(image_array, axis=0) # Add batch dimension
return image_array
@app.post('/predict')
async def predict(file: UploadFile = File(...)):
if file.content_type.startswith('image/') is False:
raise HTTPException(status_code=400, detail='Invalid file type')
image_bytes = await file.read()
image_array = preprocess_image(image_bytes)
# Make prediction
predictions = model.predict(image_array)
predicted_class = np.argmax(predictions, axis=1)[0]
# Map to category names
predicted_label = categories.get(str(predicted_class), 'Unknown')
return JSONResponse(content={
'class': predicted_label,
'confidence': float(predictions[0][predicted_class])
})
if __name__ == '__main__':
import uvicorn
uvicorn.run(app, host='0.0.0.0', port=8080)