Spaces:
Sleeping
Sleeping
| 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" | |
| 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)}) | |