| """ |
| 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 |
|
|
| |
|
|
| 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] |
|
|
| |
| CONF_THRESHOLD = 0.40 |
| TOP2_GAP_THRESHOLD = 0.15 |
| NOT_LEAF_CLASS = "not_a_leaf" |
|
|
| USE_TTA = True |
|
|
|
|
| |
|
|
| 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" |
| ) |
|
|
| |
| with open(CLASSES_PATH) as f: |
| self.classes = [x.strip() for x in f if x.strip()] |
|
|
| self.num_classes = len(self.classes) |
|
|
| |
| 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() |
|
|
| |
| if DISEASE_INFO_PATH.exists(): |
| with open(DISEASE_INFO_PATH) as f: |
| self.disease_info = json.load(f) |
| else: |
| self.disease_info = {} |
|
|
| |
| self.transform = transforms.Compose([ |
| transforms.Resize((RESIZE_TO, RESIZE_TO)), |
| transforms.CenterCrop(IMG_SIZE), |
| transforms.ToTensor(), |
| transforms.Normalize(MEAN, STD), |
| ]) |
|
|
| |
| 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.") |
|
|
| |
|
|
| 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") |
|
|
| |
|
|
| @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] |
|
|
| |
|
|
| 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) |
|
|
| |
| 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]) |
|
|
| |
| second_prob = float(probs[top5_idx[1]]) |
| gap = confidence - second_prob |
|
|
| |
| |
| |
|
|
| |
| if confidence > CONF_THRESHOLD and gap > TOP2_GAP_THRESHOLD: |
| is_leaf = True |
|
|
| |
| elif confidence > 0.30: |
| is_leaf = True |
|
|
| |
| else: |
| return self._not_leaf(top5, probs, confidence) |
|
|
| |
| |
| |
|
|
| 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, |
| } |
|
|
| |
|
|
| 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 |
| } |
