predict.py

"""
predict.py — FIXED (PRODUCTION VERSION)

Major fixes:
1. Removed over-strict rejection logic
2. Lowered confidence threshold (0.65 → 0.40)
3. Top-2 gap based decision (more reliable)
4. Reduced TTA (6 → 3 transforms)
5. Never reject obvious leaves
6. Better handling of low-confidence predictions
"""

import json
import urllib.request
from io import BytesIO
from pathlib import Path
from typing import Dict, List, Union

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

from model import build_model

# ─── CONFIG ───────────────────────────────────────────────────────────────

MODEL_PATH = Path("models/best_model.pth")
CLASSES_PATH = Path("data/classes.txt")
DISEASE_INFO_PATH = Path("data/disease_info.json")

IMG_SIZE = 300
RESIZE_TO = 332

MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]

# 🔥 FIXED THRESHOLDS
CONF_THRESHOLD = 0.40       # was 0.65 ❌
TOP2_GAP_THRESHOLD = 0.15   # was 0.25 ❌
NOT_LEAF_CLASS = "not_a_leaf"

USE_TTA = True


# ─── MODEL ────────────────────────────────────────────────────────────────

class LeafDiseasePredictor:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
        self._load()

    def _load(self):
        print("Loading model...")

        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        # Load classes
        with open(CLASSES_PATH) as f:
            self.classes = [x.strip() for x in f if x.strip()]

        self.num_classes = len(self.classes)

        # Load model
        self.model = build_model(self.num_classes, pretrained=False)
        ckpt = torch.load(MODEL_PATH, map_location=self.device)
        self.model.load_state_dict(ckpt["model_state"])
        self.model.to(self.device)
        self.model.eval()

        # Disease info
        if DISEASE_INFO_PATH.exists():
            with open(DISEASE_INFO_PATH) as f:
                self.disease_info = json.load(f)
        else:
            self.disease_info = {}

        # Transform (correct)
        self.transform = transforms.Compose([
            transforms.Resize((RESIZE_TO, RESIZE_TO)),
            transforms.CenterCrop(IMG_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(MEAN, STD),
        ])

        # 🔥 REDUCED TTA (3 instead of 6)
        self.tta_transforms = [
            self.transform,
            transforms.Compose([
                transforms.Resize((RESIZE_TO, RESIZE_TO)),
                transforms.CenterCrop(IMG_SIZE),
                transforms.RandomHorizontalFlip(p=1.0),
                transforms.ToTensor(),
                transforms.Normalize(MEAN, STD),
            ]),
            transforms.Compose([
                transforms.Resize((RESIZE_TO, RESIZE_TO)),
                transforms.RandomCrop(IMG_SIZE),
                transforms.ToTensor(),
                transforms.Normalize(MEAN, STD),
            ]),
        ]

        print("Model ready.")

    # ─── IMAGE LOADING ─────────────────────────────────────────────────────

    def _load_image(self, source):
        if isinstance(source, Image.Image):
            return source.convert("RGB")

        if isinstance(source, np.ndarray):
            return Image.fromarray(source).convert("RGB")

        source = str(source)

        if source.startswith("http"):
            with urllib.request.urlopen(source) as r:
                return Image.open(BytesIO(r.read())).convert("RGB")

        return Image.open(source).convert("RGB")

    # ─── PREDICTION CORE ───────────────────────────────────────────────────

    @torch.no_grad()
    def _predict_probs(self, img):
        probs_all = []

        if USE_TTA:
            for tf in self.tta_transforms:
                x = tf(img).unsqueeze(0).to(self.device)
                logits = self.model(x)
                probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
                probs_all.append(probs)

            return np.mean(probs_all, axis=0)

        else:
            x = self.transform(img).unsqueeze(0).to(self.device)
            logits = self.model(x)
            return F.softmax(logits, dim=-1).cpu().numpy()[0]

    # ─── MAIN PREDICT ──────────────────────────────────────────────────────

    def predict(self, source) -> Dict:
        try:
            img = self._load_image(source)
        except Exception as e:
            return self._error(f"Invalid image: {e}")

        probs = self._predict_probs(img)

        # Top-5
        top5_idx = probs.argsort()[::-1][:5]
        top5 = [
            {"class": self.classes[i], "probability": float(probs[i])}
            for i in top5_idx
        ]

        pred_idx = int(probs.argmax())
        pred_cls = self.classes[pred_idx]
        confidence = float(probs[pred_idx])

        # Top-2 gap
        second_prob = float(probs[top5_idx[1]])
        gap = confidence - second_prob

        # ─────────────────────────────────────────
        # 🔥 NEW DECISION LOGIC (CORE FIX)
        # ─────────────────────────────────────────

        # Case 1: VERY CLEAR prediction → accept
        if confidence > CONF_THRESHOLD and gap > TOP2_GAP_THRESHOLD:
            is_leaf = True

        # Case 2: Medium confidence but still reasonable → accept with warning
        elif confidence > 0.30:
            is_leaf = True

        # Case 3: Very low confidence → only then reject
        else:
            return self._not_leaf(top5, probs, confidence)

        # ─────────────────────────────────────────
        # Parse result
        # ─────────────────────────────────────────

        parts = pred_cls.split("___")
        plant = parts[0].replace("_", " ")
        disease = parts[1].replace("_", " ") if len(parts) > 1 else "Unknown"

        info = self.disease_info.get(pred_cls, {})

        warning = None
        if confidence < 0.50:
            warning = "Low confidence — try another image for confirmation."

        return {
            "is_leaf": is_leaf,
            "predicted_class": pred_cls,
            "plant": plant,
            "disease": disease,
            "confidence": confidence,
            "confidence_pct": f"{confidence:.1%}",
            "severity": info.get("severity", "Unknown"),
            "description": info.get("description", ""),
            "treatment": info.get("treatment", ""),
            "top5": top5,
            "warning": warning,
        }

    # ─── HELPERS ────────────────────────────────────────────────────────────

    def _not_leaf(self, top5, probs, confidence):
        return {
            "is_leaf": False,
            "predicted_class": NOT_LEAF_CLASS,
            "plant": "N/A",
            "disease": "N/A",
            "confidence": confidence,
            "confidence_pct": f"{confidence:.1%}",
            "severity": "N/A",
            "description": "Image not recognized as a leaf.",
            "treatment": "Upload a clear leaf image.",
            "top5": top5,
            "warning": "Model is unsure — likely not a valid leaf image.",
        }

    def _error(self, msg):
        return {
            "is_leaf": False,
            "error": msg
        }