api.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import torch
from torchvision import transforms
from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights
from PIL import Image
import io
import os

app = FastAPI(title="AgroVision API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

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

# Load models safely
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
VERIFY_MODEL_PATH = os.path.join(SCRIPT_DIR, "Jverify_model.pth")
DISEASE_MODEL_PATH = os.path.join(SCRIPT_DIR, "Jmodel.pth")

# Global dict for caching
MODELS = {
    "verify": None,
    "verify_classes": [],
    "disease": None,
    "disease_classes": []
}

def load_model(path):
    if not os.path.exists(path):
        return None, []
    checkpoint = torch.load(path, map_location=device)
    classes = checkpoint['classes']
    model = mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V1)
    model.classifier[3] = torch.nn.Linear(model.classifier[3].in_features, len(classes))
    model.load_state_dict(checkpoint['mobilenet'])
    model = model.to(device)
    model.eval()
    return model, classes

@app.on_event("startup")
def startup_event():
    print("Loading models...")
    MODELS["verify"], MODELS["verify_classes"] = load_model(VERIFY_MODEL_PATH)
    MODELS["disease"], MODELS["disease_classes"] = load_model(DISEASE_MODEL_PATH)
    print("Models loaded successfully!")

test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

@app.post("/analyze")
async def analyze_image(file: UploadFile = File(...)):
    if MODELS["verify"] is None or MODELS["disease"] is None:
        raise HTTPException(status_code=500, detail="Models not loaded")

    try:
        contents = await file.read()
        if not contents:
            raise ValueError("The server received an empty file (0 bytes). This is usually caused by a proxy stripping the file.")
            
        try:
            image = Image.open(io.BytesIO(contents)).convert('RGB')
        except Exception as pil_err:
            raise ValueError(f"Server received {len(contents)} bytes, but PIL failed to read it as an image. Details: {str(pil_err)}")
            
        image_tensor = test_transform(image).unsqueeze(0).to(device)
    except Exception as e:
        print(f"FAILED TO LOAD IMAGE: {e}")
        raise HTTPException(status_code=400, detail=str(e))

    with torch.no_grad():
        # First verify if it's a tomato leaf
        v_out = MODELS["verify"](image_tensor)
        v_probs = torch.softmax(v_out, dim=1).cpu().numpy()[0]
        v_idx = v_probs.argmax()
        is_tomato = MODELS["verify_classes"][v_idx] == "Tomatoes"
        
        if not is_tomato:
            return {
                "success": True,
                "is_tomato": False,
                "confidence": float(v_probs[v_idx]),
                "message": "This doesn't seem to be a tomato leaf. Please upload a clear image of a tomato leaf for best results.",
                "suggestion": "Make sure the leaf fills most of the frame, is well-lit, and in focus."
            }

        # It's a tomato, check for diseases
        d_out = MODELS["disease"](image_tensor)
        d_probs = torch.softmax(d_out, dim=1).cpu().numpy()[0]
        d_idx = d_probs.argmax()
        disease_name = MODELS["disease_classes"][d_idx]
        d_confidence = float(d_probs[d_idx])

        # Prepare highly professional insights based on common disease labels
        insights = {
            "Late blight": (
                "Overview: A highly destructive fungal disease (Phytophthora infestans) that thrives in cool, wet weather. "
                "Symptoms: Dark, water-soaked spots on leaves that rapidly enlarge, white fungal growth on undersides.\n\n"
                "🛡️ Treatment & Solutions:\n"
                "• Immediate Action: Entirely remove and destroy infected plant parts immediately. Do not compost.\n"
                "• Chemical Control: Apply fungicides containing chlorothalonil, copper fungicide, or mancozeb as scheduled sprays.\n"
                "• Preventative Care: Maintain excellent air circulation; water at the base of the plant avoiding foliage; space plants properly."
            ),
            "Early blight": (
                "Overview: A common fungal disease (Alternaria solani) appearing on older foliage first.\n"
                "Symptoms: Brown or black spots with characteristic concentric rings (target-board appearance) surrounded by yellow halos.\n\n"
                "🛡️ Treatment & Solutions:\n"
                "• Organic Control: Apply copper-based or sulfur-based fungicides organically approved. Neem oil can be used as a preventative measure.\n"
                "• Chemical Control: Use chlorothalonil or azoxystrobin early in the infection cycle.\n"
                "• Preventative Care: Crop rotation (3-4 years without nightshades), mulch to prevent soil splashing, and remove bottom leaves touching the soil."
            ),
            "Bacterial spot": (
                "Overview: Caused by Xanthomonas bacteria, severely impacting yield in warm, humid conditions.\n"
                "Symptoms: Small, dark, water-soaked, scabby spots on fruit and leaves.\n\n"
                "🛡️ Treatment & Solutions:\n"
                "• Immediate Action: Bacteria are notoriously difficult to control once established. Remove infected crop debris post-harvest.\n"
                "• Chemical Control: Preventative copper fungicide sprays combined with mancozeb can suppress the spread.\n"
                "• Preventative Care: Use certified disease-free seeds/transplants; avoid working among wet plants to prevent spreading bacteria on hands and tools."
            ),
            "Spider Mites": (
                "Overview: Tiny arachnids that congregate on the undersides of leaves, thriving in hot, dry conditions.\n"
                "Symptoms: Stippling (tiny yellow or white dots on leaves), fine webbing, and eventual leaf yellowing/dropping.\n\n"
                "🛡️ Treatment & Solutions:\n"
                "• Organic Control: Knock them off with a strong blast of water. Apply insecticidal soap, Neem oil, or introduce predatory mites (Phytoseiulus persimilis).\n"
                "• Chemical Control: Use targeted miticides (abamectin or spiromesifen) if the infestation is severe.\n"
                "• Preventative Care: Keep plants well-watered (stress increases susceptibility) and maintain higher localized humidity."
            ),
            "Healthy": (
                "Overview: Your tomato plant exhibits strong, vigorous growth with no current signs of widespread pathogens or pests.\n\n"
                "🌱 Ongoing Best Practices:\n"
                "• Nutrition: Ensure a balanced supply of Phosphorus and Potassium during fruiting.\n"
                "• Watering: Maintain consistent soil moisture (avoiding wet-dry extremes to prevent blossom end rot).\n"
                "• Maintenance: Continue pruning suckers for indeterminate varieties to maximize airflow and sunlight penetration."
            )
        }
        
        suggestion = insights.get(disease_name, "Consult a local agricultural expert for treatment options tailored to your region.")

        return {
            "success": True,
            "is_tomato": True,
            "prediction": disease_name,
            "confidence": d_confidence,
            "suggestion": suggestion,
            "all_probabilities": {str(cls): float(prob) for cls, prob in zip(MODELS["disease_classes"], d_probs)}
        }

if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)