PotatoGuardAPI / app.py
Liantsoaxx08's picture
add LLM for recommendation
2c67191
Raw
History Blame Contribute Delete
3.84 kB
from fastapi import FastAPI, File, UploadFile, HTTPException
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
import io
from fastapi.middleware.cors import CORSMiddleware
from llm_client import LLMClient
import json
# Initialize FastAPI app
app = FastAPI(title="Image Classification API")
import logging
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Load the Keras model once at startup
try:
model = load_model('IAPLD.h5')
except Exception as e:
raise RuntimeError(f"Failed to load model 'IAPLD.h5': {str(e)}")
# Define class names (adjust if model outputs 3 classes instead of 4)
CLASS_NAMES = ['Potato___healthy', 'Potato___Early_blight','Potato___Late_blight']
# Function to preprocess the uploaded image
def preprocess_image(image: Image.Image) -> np.ndarray:
# Resize to match model input shape (250, 250 as per your code)
image = image.resize((250, 250)) # Adjust to (256, 256) if model expects that
# Convert to NumPy array and normalize to 0-1 range
image_array = np.array(image) / 255.0
# Add batch dimension (1, 250, 250, 3)
image_array = np.expand_dims(image_array, axis=0)
return image_array
# Root endpoint
@app.get("/")
async def root():
return {"message": "Welcome to the Image Classification API. Use POST /predict/ to upload an image."}
from recomm import Redommend
@app.get("/recommendation")
async def recommendation(disease: str):
if not disease:
raise HTTPException(status_code=400, detail="Disease parameter is required")
try:
llm_client = LLMClient()
recommender = Redommend(llm_client)
raw_response = recommender._run(disease).strip()
if raw_response.startswith("```json"):
raw_response = raw_response.replace("```json", "").replace("```", "").strip()
data = json.loads(raw_response)
return data
except json.JSONDecodeError:
raise HTTPException(
status_code=500,
detail=f"Le LLM n’a pas renvoyé un JSON valide : {raw_response}"
)
except Exception as e:
logging.error(f"Error in recommendation endpoint: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")
# Prediction endpoint
@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
if not file.content_type.startswith('image/'):
raise HTTPException(status_code=400, detail="Uploaded file must be an image")
try:
# Read the image bytes
contents = await file.read()
# Open as PIL image
image = Image.open(io.BytesIO(contents))
print("Image size:", image.size) # Debug: Check image size
# Preprocess the image
image_array = preprocess_image(image)
print("Image shape:", image_array.shape) # Debug: Check input shape
# Make prediction (model outputs probabilities directly)
predictions = model.predict(image_array)
print("Probabilities:", predictions) # Debug: Direct probabilities
# Get predicted class and confidence
class_index = np.argmax(predictions[0])
class_name = CLASS_NAMES[class_index]
probability = float(predictions[0][class_index])
# Return prediction result
return {
"predicted_class": class_name,
"confidence": probability
}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
# Run the app with: uvicorn main:app --reload
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)