Spaces:
Runtime error
Runtime error
File size: 3,108 Bytes
ca8798e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import io
from models import ResNet9
app = FastAPI(title="CropGuard - Plant Disease Detection")
CLASS_NAMES = [
'Apple___Apple_scab', 'Apple___Black_rot', 'Apple___Cedar_apple_rust', 'Apple___healthy',
'Blueberry___healthy', 'Cherry_(including_sour)__Powdery_mildew', 'Cherry(including_sour)__healthy',
'Corn(maize)__Cercospora_leaf_spot Gray_leaf_spot', 'Corn(maize)_Common_rust',
'Corn(maize)__Northern_Leaf_Blight', 'Corn(maize)healthy', 'Grape___Black_rot',
'Grape___Esca(Black_Measles)', 'Grape___Leaf_blight(Isariopsis_Leaf_Spot)', 'Grape___healthy',
'Orange___Haunglongbing(Citrus_greening)', 'Peach___Bacterial_spot', 'Peach___healthy',
'Pepper,_bell___Bacterial_spot', 'Pepper,_bell___healthy', 'Potato___Early_blight',
'Potato___Late_blight', 'Potato___healthy', 'Raspberry___healthy', 'Soybean___healthy',
'Squash___Powdery_mildew', 'Strawberry___Leaf_scorch', 'Strawberry___healthy',
'Tomato___Bacterial_spot', 'Tomato___Early_blight', 'Tomato___Late_blight', 'Tomato___Leaf_Mold',
'Tomato___Septoria_leaf_spot', 'Tomato___Spider_mites Two-spotted_spider_mite', 'Tomato___Target_Spot',
'Tomato___Tomato_Yellow_Leaf_Curl_Virus', 'Tomato___Tomato_mosaic_virus', 'Tomato___healthy'
]
model = None
def load_model():
global model
if model is None:
model = ResNet9(3, len(CLASS_NAMES))
state_dict = torch.load("plant-disease-model-state-dict.pth", map_location="cpu")
model.load_state_dict(state_dict)
model.eval()
load_model()
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
try:
image_bytes = await file.read()
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor()
])
img_tensor = transform(image)
if isinstance(img_tensor, torch.Tensor) and img_tensor.ndimension() == 3:
img_tensor = img_tensor.unsqueeze(0)
global model
if model is None:
load_model()
if model is None:
raise RuntimeError("Model failed to load.")
with torch.no_grad():
outputs = model(img_tensor)
probabilities = F.softmax(outputs[0], dim=0)
top5_prob, top5_indices = torch.topk(probabilities, 5)
results = {}
for prob, idx in zip(top5_prob, top5_indices):
class_name = CLASS_NAMES[int(idx.item())]
clean_name = class_name.replace('___', ' - ').replace('_', ' ')
results[clean_name] = float(prob)
return JSONResponse(content={"predictions": results})
except Exception as e:
return JSONResponse(content={"error": str(e)}, status_code=500)
@app.get("/")
def root():
return {"message": "CropGuard FastAPI is running. Use /predict to POST an image."}
|