Blame for api.py

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