Spaces:
Runtime error
Runtime error
File size: 1,535 Bytes
f9e98b8 |
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 |
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)
|