Update app.py

app.py

@@ -3,9 +3,9 @@ from fastapi.middleware.cors import CORSMiddleware
 from ultralytics import YOLO
 from PIL import Image
 import io
- 
+
 app = FastAPI(title="AgriVision API")
- 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -132,7 +132,7 @@ TREATMENTS = {
         "treatment": "Apply pyrethroid insecticides. Store harvested grain in airtight containers. Use hermetic storage bags and apply diatomaceous earth as a physical control."
     },
 }
- 
+
 # ── Severity function ──────────────────────────────────────────
 def get_severity(confidence: float, crop_type: str) -> str:
     if crop_type == "Healthy":
@@ -143,56 +143,74 @@ def get_severity(confidence: float, crop_type: str) -> str:
         return "Medium"
     else:
         return "Low"
- 
+
 # ── Name formatter ─────────────────────────────────────────────
 def format_class_name(name: str) -> str:
     return name.replace("_", " ").title()
- 
+
 # ── Load model once at startup ─────────────────────────────────
 model = YOLO("best.pt")
- 
+
 # ── Routes ─────────────────────────────────────────────────────
 @app.get("/")
 def home():
     return {"status": "AgriVision API is running"}
- 
+
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # Read and process image
     contents = await file.read()
     image = Image.open(io.BytesIO(contents)).convert("RGB")
- 
-    # Run inference
+
     results = model.predict(image, imgsz=224, verbose=False)
     probs   = results[0].probs
- 
-    # Top prediction
+
     top1_index      = probs.top1
     top1_class      = model.names[top1_index]
     top1_confidence = round(float(probs.top1conf), 4)
- 
-    # ── Invalid image detection ─────────────────────────────────
-    top3_probs  = [float(probs.data[i]) for i in probs.top5[:3]]
-    top3_spread = max(top3_probs) - min(top3_probs)
- 
-    if top1_confidence < 0.45 and top3_spread < 0.45:
+
+    # ── Invalid image detection ────────────────────────────────
+    # Get all top 5 probabilities
+    top5_probs = [float(probs.data[i]) for i in probs.top5]
+
+    # How much the top prediction dominates over the rest
+    top1_prob  = top5_probs[0]
+    top2_prob  = top5_probs[1]
+    top3_prob  = top5_probs[2]
+
+    # The gap between 1st and 2nd — a confident real crop image
+    # will have a clear winner. A random image spreads probability evenly.
+    top1_top2_gap = top1_prob - top2_prob
+
+    # Entropy-like spread across top 3
+    top3_spread = top1_prob - top3_prob
+
+    # Reject if EITHER of these is true:
+    # 1. Very low confidence (model has no idea)
+    # 2. Low confidence AND the top predictions are bunched together
+    #    (model is guessing randomly across classes)
+    is_invalid = (
+        top1_confidence < 0.30                          # extremely low confidence
+        or (top1_confidence < 0.55 and top1_top2_gap < 0.10)  # low + no clear winner
+        or (top1_confidence < 0.50 and top3_spread < 0.15)    # low + very spread out
+    )
+
+    if is_invalid:
         raise HTTPException(
             status_code=422,
             detail={
                 "error": "Invalid image",
-                "message": "The uploaded image does not appear to be a crop or pest image. Please upload a clear photo of a crop leaf or pest.",
+                "message": "The uploaded image does not appear to be a crop or pest. Please upload a clear photo of a crop leaf or pest.",
                 "confidence": f"{round(top1_confidence * 100, 2)}%"
             }
         )
- 
-    # Look up treatment and type
+
+    # ── Valid scan ─────────────────────────────────────────────
     info     = TREATMENTS.get(top1_class, {
         "type": "Unknown",
         "treatment": "No treatment data available for this class."
     })
     severity = get_severity(top1_confidence, info["type"])
- 
-    # Top 3 predictions
+
     top3 = [
         {
             "class":      format_class_name(model.names[i]),
@@ -200,7 +218,7 @@ async def predict(file: UploadFile = File(...)):
         }
         for i in probs.top5[:3]
     ]
- 
+
     return {
         "prediction":   format_class_name(top1_class),
         "type":         info["type"],