Spaces:
Sleeping
Sleeping
| 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 | |
| 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]) | |
| ]) | |
| 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) | |