Spaces:
Sleeping
Sleeping
| 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() | |
| def health() -> dict: | |
| return {"status": "ok", "classes": len(CLASS_NAMES)} | |
| 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 | |