Update main.py

main.py

@@ -1,11 +1,10 @@
 from fastapi import FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
-from ultralytics import YOLO  # Import the AI library
+from sahi import AutoDetectionModel
+from sahi.predict import get_sliced_prediction
 import uvicorn
 import shutil
 import os
-import cv2
-import numpy as np
 
 app = FastAPI()
 
@@ -18,18 +17,25 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# LOAD THE MODEL ONCE (At startup)
-# Make sure 'best.pt' is in the same folder as main.py!
-print("⏳ Loading AI Model...")
+# ==========================================
+# 🧠 LOAD SAHI MODEL
+# ==========================================
+print("⏳ Loading SAHI + YOLO Model...")
 try:
-    model = YOLO("best.pt")
-    print("✅ Model loaded successfully!")
+    # We wrap your YOLO model inside SAHI
+    detection_model = AutoDetectionModel.from_pretrained(
+        model_type='yolov8',
+        model_path='best.pt',    # Your trained model
+        confidence_threshold=0.25, # Lower threshold because SAHI is more accurate
+        device='cpu'             # Use 'cpu' for Hugging Face free tier
+    )
+    print("✅ SAHI Model loaded successfully!")
 except Exception as e:
     print(f"❌ Error loading model: {e}")
 
 @app.get("/")
 def home():
-    return {"message": "Goyam AI Server is Running! 🚀"}
+    return {"message": "Goyam AI (SAHI Enabled) is Running! 🚀"}
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
@@ -41,29 +47,31 @@ async def predict(file: UploadFile = File(...)):
         shutil.copyfileobj(file.file, buffer)
 
     try:
-        # 2. RUN REAL AI PREDICTION
-        results = model(temp_filename)
+        # 2. RUN SAHI SLICED PREDICTION
+        # This cuts the image into 512x512 blocks and checks each one
+        result = get_sliced_prediction(
+            temp_filename,
+            detection_model,
+            slice_height=512,
+            slice_width=512,
+            overlap_height_ratio=0.2,
+            overlap_width_ratio=0.2
+        )
         
         # 3. Process Results
-        # Get the first result (since we sent one image)
-        result = results[0]
+        predictions = result.object_prediction_list
         
-        # Check if any objects were detected
-        if len(result.boxes) > 0:
-            # Find the box with the highest confidence
-            # box.conf is a list of confidences, we take the max
-            best_box = max(result.boxes, key=lambda x: x.conf[0])
+        if len(predictions) > 0:
+            # Find the prediction with the highest confidence score
+            best_pred = max(predictions, key=lambda x: x.score.value)
             
-            class_id = int(best_box.cls[0])
-            confidence_score = float(best_box.conf[0])
-            
-            # Get the name from the model's internal names list
-            detected_name = result.names[class_id]
+            detected_name = best_pred.category.name
+            confidence_score = best_pred.score.value
             
             # Format confidence as percentage
             confidence_percent = f"{int(confidence_score * 100)}%"
             
-            print(f"✅ Detected: {detected_name} ({confidence_percent})")
+            print(f"✅ Detected with SAHI: {detected_name} ({confidence_percent})")
             
             return {
                 "filename": file.filename,
@@ -72,12 +80,12 @@ async def predict(file: UploadFile = File(...)):
                 "recommendation": get_recommendation(detected_name)
             }
         else:
-            print("⚠️ No disease detected")
+            print("⚠️ No disease detected (SAHI)")
             return {
                 "filename": file.filename,
-                "disease": "No Disease Detected",
+                "disease": "Healthy / No Detection",
                 "confidence": "0%",
-                "recommendation": "Plant looks healthy or unclear. Try closer photo."
+                "recommendation": "Plant looks healthy. Ensure regular water supply."
             }
 
     except Exception as e:
@@ -85,11 +93,11 @@ async def predict(file: UploadFile = File(...)):
         return {"error": str(e)}
         
     finally:
-        # Clean up (Delete the temp file) regardless of success/fail
+        # Clean up
         if os.path.exists(temp_filename):
             os.remove(temp_filename)
 
-# Helper function to give advice based on disease
+# Helper function
 def get_recommendation(disease_name):
     recommendations = {
         "Blast": "Use Tricyclazole 75 WP. Avoid excess nitrogen.",
@@ -101,4 +109,4 @@ def get_recommendation(disease_name):
     return recommendations.get(disease_name, "Consult an agricultural officer.")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=7860)