Spaces:
Sleeping
Sleeping
from fastapi import FastAPI, File, UploadFile, HTTPException | |
from fastapi.responses import JSONResponse | |
from tensorflow.keras.models import load_model | |
from tensorflow.keras.preprocessing import image | |
import numpy as np | |
import logging | |
from PIL import Image | |
import io | |
# Configure logging | |
logging.basicConfig(level=logging.DEBUG) | |
# Initialize FastAPI app | |
app = FastAPI() | |
# Load your trained model | |
model = load_model('model.h5') | |
class_names = ['Normal', 'bacteria', 'virus'] | |
def preprocess_image(img, target_size): | |
"""Resize and preprocess the image for the model.""" | |
if img.mode != "RGB": | |
img = img.convert("RGB") | |
img = img.resize(target_size) | |
img_array = image.img_to_array(img) | |
img_array = np.expand_dims(img_array, axis=0) # Add batch dimension | |
return img_array | |
async def predict(file: UploadFile = File(...)): | |
if not file: | |
raise HTTPException(status_code=400, detail="No file provided") | |
try: | |
# Read the file's content into a BytesIO object | |
img_bytes = io.BytesIO(await file.read()) | |
# Use PIL to open the image | |
img = Image.open(img_bytes) | |
img_array = preprocess_image(img, (224, 224)) | |
# Make prediction | |
predictions = model.predict(img_array) | |
predicted_class = np.argmax(predictions, axis=1) | |
# Return the prediction | |
predictions = { | |
'class': class_names[predicted_class[0]], | |
'confidence': float(predictions[0][predicted_class[0]]) | |
} | |
return JSONResponse(content=predictions) | |
except Exception as e: | |
logging.debug(f"Error processing the file: {str(e)}") | |
raise HTTPException(status_code=500, detail=f"Error processing the file: {str(e)}") | |
if __name__ == '__main__': | |
app.run() |