from pathlib import Path from io import BytesIO import numpy as np import tensorflow as tf from fastapi import FastAPI, File, HTTPException, UploadFile from fastapi.middleware.cors import CORSMiddleware from PIL import Image MODEL_PATH = Path("plant_village_model.keras") LABELS_PATH = Path("PlantDiseaseDetectorExpo") / "assets" / "labels.txt" TFLITE_PATH = Path("PlantDiseaseDetectorExpo") / "assets" / "plant_model.tflite" DEFAULT_CLASS_NAMES = [ "Apple___Apple_scab", "Apple___Black_rot", "Apple___Cedar_apple_rust", "Apple___healthy", "Blueberry___healthy", "Cherry___healthy", "Cherry___Powdery_mildew", "Corn___Cercospora_leaf_spot Gray_leaf_spot", "Corn___Common_rust", "Corn___healthy", "Corn___Northern_Leaf_Blight", "Grape___Black_rot", "Grape___Esca_(Black_Measles)", "Grape___healthy", "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)", "Orange___Haunglongbing_(Citrus_greening)", "Peach___Bacterial_spot", "Peach___healthy", "Pepper,_bell___Bacterial_spot", "Pepper,_bell___healthy", "Potato___Early_blight", "Potato___healthy", "Potato___Late_blight", "Raspberry___healthy", "Soybean___healthy", "Squash___Powdery_mildew", "Strawberry___healthy", "Strawberry___Leaf_scorch", "Tomato___Bacterial_spot", "Tomato___Early_blight", "Tomato___healthy", "Tomato___Late_blight", "Tomato___Leaf_Mold", "Tomato___Septoria_leaf_spot", "Tomato___Spider_mites Two-spotted_spider_mite", "Tomato___Target_Spot", "Tomato___Tomato_mosaic_virus", "Tomato___Tomato_Yellow_Leaf_Curl_Virus", ] app = FastAPI(title="Plant Disease Detection API", version="1.0.0") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class DenseCompat(tf.keras.layers.Dense): """Compat layer that ignores unknown Dense config keys from newer Keras saves.""" def __init__(self, *args, quantization_config=None, **kwargs): super().__init__(*args, **kwargs) class TFLiteModelWrapper: """Wrapper to expose a Keras-like predict() around a TFLite interpreter.""" def __init__(self, model_path: Path): self.interpreter = tf.lite.Interpreter(model_path=str(model_path)) self.interpreter.allocate_tensors() self.input_details = self.interpreter.get_input_details()[0] self.output_details = self.interpreter.get_output_details()[0] def predict(self, batch: np.ndarray) -> np.ndarray: input_data = np.asarray(batch, dtype=np.float32) self.interpreter.set_tensor(self.input_details["index"], input_data) self.interpreter.invoke() output = self.interpreter.get_tensor(self.output_details["index"]) return np.asarray(output) def load_class_names() -> list[str]: if not LABELS_PATH.exists(): return DEFAULT_CLASS_NAMES.copy() class_names = [] for line in LABELS_PATH.read_text(encoding="utf-8").splitlines(): line = line.strip() if not line: continue if ":" in line: class_names.append(line.split(":", 1)[1].strip()) else: class_names.append(line) return class_names def load_model(): custom_objects = { "Dense": DenseCompat, "keras.layers.Dense": DenseCompat, "keras.src.layers.core.dense.Dense": DenseCompat, } try: return tf.keras.models.load_model(MODEL_PATH, compile=False) except TypeError as exc: if "quantization_config" in str(exc): try: return tf.keras.models.load_model( MODEL_PATH, compile=False, custom_objects=custom_objects, ) except Exception: pass if TFLITE_PATH.exists(): return TFLiteModelWrapper(TFLITE_PATH) raise RuntimeError( "Could not load .keras model and no TFLite fallback was found at " f"{TFLITE_PATH}" ) def preprocess_image(image_bytes: bytes) -> np.ndarray: image = Image.open(BytesIO(image_bytes)).convert("RGB") image = image.resize((224, 224)) rgb_uint8 = np.array(image, dtype=np.uint8) img_array = rgb_uint8.astype(np.float32) / 255.0 return np.expand_dims(img_array, axis=0), rgb_uint8 def should_reject_as_no_leaf(rgb_uint8: np.ndarray, probs: np.ndarray) -> tuple[bool, str]: confidence = float(np.max(probs)) sorted_probs = np.sort(probs) second_best = float(sorted_probs[-2]) if sorted_probs.size > 1 else 0.0 margin = confidence - second_best r = rgb_uint8[:, :, 0].astype(np.float32) g = rgb_uint8[:, :, 1].astype(np.float32) b = rgb_uint8[:, :, 2].astype(np.float32) green_mask = (g > 40) & (g > 1.05 * r) & (g > 1.05 * b) green_ratio = float(np.mean(green_mask)) if confidence < 0.55: return True, "low_confidence" if margin < 0.12 and confidence < 0.70: return True, "ambiguous_prediction" if green_ratio < 0.05 and confidence < 0.80: return True, "non_leaf_visual_pattern" return False, "accepted" MODEL = load_model() CLASS_NAMES = load_class_names() @app.get("/health") def health() -> dict: return {"status": "ok", "classes": len(CLASS_NAMES)} @app.post("/predict") async def predict(file: UploadFile = File(...)) -> dict: if not file.content_type or not file.content_type.startswith("image/"): raise HTTPException(status_code=400, detail="Please upload a valid image file") try: image_bytes = await file.read() input_tensor, rgb_uint8 = preprocess_image(image_bytes) preds = MODEL.predict(input_tensor) probs = np.asarray(preds[0], dtype=np.float32) result_idx = int(np.argmax(probs)) confidence = float(probs[result_idx]) if 0 <= result_idx < len(CLASS_NAMES): label = CLASS_NAMES[result_idx] else: label = f"Class {result_idx}" top3_idx = np.argsort(probs)[-3:][::-1] top3 = [] for idx in top3_idx: class_name = CLASS_NAMES[idx] if idx < len(CLASS_NAMES) else f"Class {int(idx)}" top3.append({ "class_index": int(idx), "class_name": class_name, "confidence": float(probs[idx]), }) reject, reject_reason = should_reject_as_no_leaf(rgb_uint8, probs) if reject: return { "prediction": { "class_index": -1, "class_name": "No leaf detected", "confidence": float(1.0 - confidence), }, "top3": top3, "rejected": True, "reject_reason": reject_reason, } return { "prediction": { "class_index": result_idx, "class_name": label, "confidence": confidence, }, "top3": top3, "rejected": False, } except HTTPException: raise except Exception as exc: raise HTTPException(status_code=500, detail=f"Prediction failed: {exc}") from exc