Upload 2 files

best.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6a4a858112a248a4ac07bc254ccc70ca2f71048ed7b7e9018fe618ce766517cf
-size 54521258
+oid sha256:7f2e775168f7ce664b7abe08dbd5a7c903f5f3487c0711fe180fa8836015c312
+size 77368442

main.py

@@ -1,9 +1,7 @@
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.middleware.cors import CORSMiddleware
-from typing import List, Optional
-from sahi import AutoDetectionModel
-from sahi.predict import get_sliced_prediction
-from ultralytics import YOLO  # 👈 Import YOLO for classification
+from typing import Optional
+from ultralytics import YOLO
 import uvicorn
 import shutil
 import os
@@ -19,80 +17,79 @@ app.add_middleware(
 )
 
 # ==========================================
-# 🧠 LOAD MODELS
+# 🧠 LOAD MODELS (Classification Only!)
 # ==========================================
 print("⏳ Loading Models...")
 
-# 1. Disease Detection Model (Your Custom Model)
 try:
-    detection_model = AutoDetectionModel.from_pretrained(
-        model_type='yolov8',
-        model_path='best.pt',
-        confidence_threshold=0.25,
-        device='cpu'             
-    )
-    print("✅ Disease Model loaded!")
+    # 1. Disease Classification Model (Your ultra-smart 96.8% brain)
+    disease_model = YOLO('best.pt')
+    print("✅ Goyam Disease Classifier loaded!")
 except Exception as e:
     print(f"❌ Error loading Disease Model: {e}")
 
-# 2. Plant Gatekeeper Model (Generic Classifier)
-# This downloads automatically the first time (only ~6MB)
 try:
-    classifier = YOLO('yolov8n-cls.pt') 
+    # 2. Plant Gatekeeper (ImageNet 1000-class classifier)
+    gatekeeper_model = YOLO('yolov8n-cls.pt') 
     print("✅ Plant Gatekeeper loaded!")
 except Exception as e:
     print(f"❌ Error loading Gatekeeper: {e}")
 
 
 # ==========================================
-# 🛡️ THE IMPROVED PLANT GATEKEEPER
+# 🛡️ THE BULLETPROOF GATEKEEPER (Allowlist)
 # ==========================================
 def is_likely_plant(image_path):
     """
-    Revised Logic: BLOCK specific non-plant objects (People, Electronics, Cars).
-    ALLOW everything else (Nature, Fields, Valleys, etc.).
+    STRICT ALLOWLIST LOGIC:
+    We check the top 5 things the AI thinks it sees. 
+    If NONE of them are related to nature, agriculture, or plants, we reject the photo.
     """
     try:
-        # Predict top class
-        results = classifier(image_path, verbose=False)
-        top_probs = results[0].probs.top1
-        top_name = results[0].names[top_probs].lower()
+        results = gatekeeper_model(image_path, verbose=False)
         
-        print(f"🧐 Gatekeeper sees: '{top_name}'")
-
-        # ⛔ BLOCK LIST: If the AI sees these, REJECT the image.
-        forbidden_keywords = [
-            'person', 'man', 'woman', 'boy', 'girl', 'human', 'face', 'groom', 'diver',
-            'computer', 'laptop', 'screen', 'monitor', 'keyboard', 'mouse', 'phone', 'cellular',
-            'car', 'vehicle', 'truck', 'bus', 'racer', 'motor', 'machine',
-            'furniture', 'table', 'chair', 'seat', 'couch', 'desk',
-            'dog', 'cat', 'animal', 'bear', 'lion' 
+        # Get the top 5 predicted classes
+        top5_indices = results[0].probs.top5
+        top5_names = [results[0].names[i].lower() for i in top5_indices]
+        
+        print(f"🧐 Gatekeeper sees: {top5_names}")
+
+        # ✅ ALLOW LIST: Botanical, agricultural, and nature terms
+        allowed_keywords = [
+            'plant', 'leaf', 'grass', 'flower', 'tree', 'fern', 'moss', 'weed', 
+            'crop', 'agriculture', 'field', 'greenhouse', 'pot', 'earth', 'soil', 
+            'vegetation', 'forest', 'valley', 'daisy', 'corn', 'acorn', 'paddy'
         ]
 
-        # Check if it matches a forbidden item
-        for bad_word in forbidden_keywords:
-            if bad_word in top_name:
-                print(f"⛔ Blocked: Detected '{top_name}' (Matched '{bad_word}')")
-                return False
+        # Check if ANY of the top 5 predictions contain our allowed keywords
+        for predicted_item in top5_names:
+            for good_word in allowed_keywords:
+                if good_word in predicted_item:
+                    print(f"✅ Passed: Gatekeeper verified plant matter ('{predicted_item}')")
+                    return True
 
-        # ✅ If it's not forbidden, we ALLOW it.
-        # This lets "Valley", "Lakeside", "Grass", "Meadow" pass through.
-        return True
+        # If the loop finishes and didn't find a single nature word, reject it!
+        print(f"🚫 Blocked: No agricultural or plant features detected.")
+        return False
 
     except Exception as e:
         print(f"⚠️ Gatekeeper Error: {e}")
-        return True # If gatekeeper fails, let it pass to be safe
+        return True # Fail-safe: let it pass if the gatekeeper crashes
 
 
 def get_recommendation(disease_name):
     recommendations = {
-        "Blast": "Use Tricyclazole 75 WP. Avoid excess nitrogen.",
-        "Blight": "Drain field. Apply fresh water. Use bleaching powder.",
-        "Brown Spot": "Improve soil fertility. Add Potassium.",
-        "Hispa": "Spray Chloropyriphos. Clip leaf tips.",
-        "Healthy": "Keep maintaining good water levels!"
+        "Leaf Blast": "Use Tricyclazole 75 WP. Avoid applying excess nitrogen fertilizer.",
+        "Sheath Blight": "Drain the field immediately. Apply validamycin or carbendazim.",
+        "Brown Spot": "Improve soil fertility. Apply potassium and phosphorus.",
+        "Healthy Rice Leaf": "No disease detected. Keep maintaining optimal water levels!"
     }
-    return recommendations.get(disease_name, "Consult an agricultural officer.")
+    # Fallback just in case class names differ slightly
+    for key, value in recommendations.items():
+        if key.lower() in disease_name.lower():
+            return value
+    return "Consult your local agricultural extension officer for treatment."
+
 
 @app.get("/")
 def home():
@@ -111,57 +108,38 @@ async def predict(
         with open(temp_filename, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
 
-        # 🛑 STEP 1: RUN THE GATEKEEPER
+        # 🛑 STEP 1: RUN THE STRICT GATEKEEPER
         if not is_likely_plant(temp_filename):
-            print("🚫 Rejected: Image does not look like a plant.")
             return {
                 "filename": file.filename,
                 "disease": "Invalid Image",
                 "confidence": "0%",
-                "recommendation": "This image does not appear to be a plant. Please upload a clear photo of your paddy field.",
+                "recommendation": "This image does not appear to be a plant or paddy field. Please upload a clear photo of a rice leaf.",
                 "latitude": float(latitude) if latitude else None,
                 "longitude": float(longitude) if longitude else None
             }
 
-        # ✅ STEP 2: RUN DISEASE DETECTION (SAHI)
-        result = get_sliced_prediction(
-            temp_filename,
-            detection_model,
-            slice_height=640,
-            slice_width=640,
-            overlap_height_ratio=0.2,
-            overlap_width_ratio=0.2
-        )
+        # ✅ STEP 2: RUN DISEASE CLASSIFICATION
+        results = disease_model(temp_filename, verbose=False)
         
-        predictions = result.object_prediction_list
+        # Extract the highest probability prediction
+        top_idx = results[0].probs.top1
+        confidence_score = float(results[0].probs.top1conf)
+        detected_name = results[0].names[top_idx]
         
         response_data = {
             "filename": file.filename,
+            "disease": detected_name,
+            "confidence": f"{int(confidence_score * 100)}%",
+            "recommendation": get_recommendation(detected_name),
             "latitude": float(latitude) if latitude else None,
             "longitude": float(longitude) if longitude else None
         }
 
-        if len(predictions) > 0:
-            best_pred = max(predictions, key=lambda x: x.score.value)
-            detected_name = best_pred.category.name
-            confidence_score = best_pred.score.value
-            
-            response_data.update({
-                "disease": detected_name,
-                "confidence": f"{int(confidence_score * 100)}%",
-                "recommendation": get_recommendation(detected_name)
-            })
-        else:
-            response_data.update({
-                "disease": "Healthy / No Detection",
-                "confidence": "0%",
-                "recommendation": get_recommendation("Healthy")
-            })
-            
         return response_data
 
     except Exception as e:
-        print(f"Error: {e}")
+        print(f"❌ API Error: {e}")
         return {"error": str(e)}
         
     finally: