|
import os |
|
import numpy as np |
|
import tensorflow as tf |
|
from fastapi import FastAPI, File, UploadFile |
|
from fastapi.responses import JSONResponse |
|
from io import BytesIO |
|
from PIL import Image |
|
from tensorflow.keras.preprocessing.image import img_to_array |
|
from tensorflow.keras.applications import resnet50 |
|
from tensorflow.keras.applications.resnet50 import preprocess_input |
|
import uvicorn |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
model_path = "model.keras" |
|
class_indices = {0: 'blight', 1: 'brown_spots'} |
|
|
|
|
|
if os.path.exists(model_path): |
|
model = tf.keras.models.load_model(model_path) |
|
print("Model loaded successfully.") |
|
else: |
|
print(f"Model file not found at {model_path}. Please upload the model.") |
|
|
|
|
|
def predict_image(image_data): |
|
try: |
|
|
|
img = Image.open(BytesIO(image_data)) |
|
|
|
img = img.resize((224, 224)) |
|
|
|
img_array = img_to_array(img) |
|
img_array = np.expand_dims(img_array, axis=0) |
|
img_array = preprocess_input(img_array) |
|
|
|
|
|
prediction = model.predict(img_array) |
|
predicted_class = np.argmax(prediction[0]) |
|
class_name = class_indices[predicted_class] |
|
return class_name |
|
except Exception as e: |
|
print("Prediction error:", e) |
|
return "Error during prediction" |
|
|
|
|
|
@app.get("/health") |
|
async def api_health_check(): |
|
return JSONResponse(content={"status": "Service is running"}) |
|
|
|
|
|
@app.post("/predict") |
|
async def api_predict_image(file: UploadFile = File(...)): |
|
try: |
|
|
|
image_data = await file.read() |
|
|
|
|
|
prediction = predict_image(image_data) |
|
|
|
return JSONResponse(content={"prediction": prediction}) |
|
except Exception as e: |
|
return JSONResponse(content={"error": str(e)}) |
|
|
|
|
|
if __name__ == "__main__": |
|
uvicorn.run(app, host="0.0.0.0", port=7860) |
|
|