| | import tensorflow as tf
|
| | from fastapi import FastAPI, File, UploadFile, HTTPException
|
| | from fastapi.responses import JSONResponse
|
| | from fastapi.middleware.cors import CORSMiddleware
|
| | from io import BytesIO
|
| | from PIL import Image
|
| | import numpy as np
|
| | from typing import List
|
| | import logging
|
| |
|
| | logging.basicConfig(level=logging.INFO)
|
| | logger = logging.getLogger(__name__)
|
| |
|
| | app = FastAPI()
|
| |
|
| | app.add_middleware(
|
| | CORSMiddleware,
|
| | allow_origins=["*"],
|
| | allow_credentials=True,
|
| | allow_methods=["*"],
|
| | allow_headers=["*"],
|
| | )
|
| |
|
| | MODEL_PATH = "model.keras"
|
| | try:
|
| | logger.info(f"Loading model from {MODEL_PATH}")
|
| | model = tf.keras.models.load_model(MODEL_PATH)
|
| | except Exception as e:
|
| | logger.error(f"Failed to load model: {e}")
|
| | raise
|
| |
|
| | CLASS_NAMES = ['healthy', 'rotten']
|
| |
|
| | def process_image(image: Image.Image) -> np.ndarray:
|
| | try:
|
| | if image.mode != 'RGB':
|
| | image = image.convert('RGB')
|
| | image = image.resize((128, 128))
|
| | image_array = np.array(image) / 255.0
|
| | if image_array.shape[-1] == 4:
|
| | image_array = image_array[..., :3]
|
| | return image_array
|
| | except Exception as e:
|
| | logger.error(f"Error processing image: {e}")
|
| | raise ValueError(f"Error processing image: {e}")
|
| |
|
| | @app.post("/predict")
|
| | async def predict(files: List[UploadFile] = File(...)):
|
| | try:
|
| | if not files:
|
| | logger.warning("No files provided in request")
|
| | raise HTTPException(status_code=400, detail="No files provided")
|
| |
|
| | results = []
|
| | for file in files:
|
| | logger.info(f"Processing file: {file.filename}")
|
| |
|
| | if file.size > 10 * 1024 * 1024:
|
| | logger.warning(f"File {file.filename} too large: {file.size} bytes")
|
| | raise HTTPException(status_code=400, detail=f"File {file.filename} too large")
|
| |
|
| |
|
| | content = await file.read()
|
| | image = Image.open(BytesIO(content))
|
| | image_array = process_image(image)
|
| | image_array = np.expand_dims(image_array, axis=0)
|
| |
|
| |
|
| | logger.info(f"Making prediction for {file.filename}")
|
| | predictions = model.predict(image_array, verbose=0)
|
| | predicted_index = np.argmax(predictions[0])
|
| | predicted_class = CLASS_NAMES[predicted_index]
|
| | confidence = float(np.max(predictions[0]))
|
| |
|
| | results.append({
|
| | "filename": file.filename or f"image_{len(results)}.jpg",
|
| | "class": predicted_class,
|
| | "confidence": round(confidence, 4)
|
| | })
|
| |
|
| | logger.info(f"Returning {len(results)} predictions")
|
| | return JSONResponse({
|
| | "predictions": results,
|
| | "status": "success"
|
| | })
|
| |
|
| | except Exception as e:
|
| | logger.error(f"Error in predict endpoint: {e}")
|
| | return JSONResponse(
|
| | status_code=400,
|
| | content={
|
| | "predictions": [],
|
| | "status": "error",
|
| | "error": str(e)
|
| | }
|
| | )
|
| |
|
| |
|
