SanchaiKB/Insect-Classification-Model

1import torch2import torch.nn as nn3from torchvision import models, transforms4from PIL import Image5from fastapi import FastAPI, UploadFile, File6from io import BytesIO7 8app = FastAPI(title="Rice Insect Classifier API")9 10# 1. Load Model11def load_model():12    model = models.efficientnet_b0()13    num_ftrs = model.classifier[1].in_features14    model.classifier = nn.Sequential(15        nn.Identity(),16        nn.Sequential(17            nn.Dropout(p=0.2, inplace=True),18            nn.Linear(num_ftrs, 5)19        )20    )21    model.load_state_dict(torch.load("insect_efficientnet.pth", map_location='cpu'))22    model.eval()23    return model24 25model = load_model()26labels = [27    "Rice stem borer",28    "green leafhopper",29    "planthopper",30    "rice bug",31    "rice leaf roller"32]33 34# 2. Endpoints35@app.get("/health")36def health():37    return {"status": "ok", "model": "efficientnet_b0"}38 39@app.post("/predict")40async def predict(file: UploadFile = File(...)):41    # Preprocessing42    transform = transforms.Compose([43        transforms.Resize((224, 224)),44        transforms.ToTensor(),45        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])46    ])47    48    image_bytes = await file.read()49    image = Image.open(BytesIO(image_bytes)).convert("RGB")50    img_t = transform(image).unsqueeze(0)51    52    # Inference53    with torch.no_grad():54        logits = model(img_t)55        probs = torch.nn.functional.softmax(logits[0], dim=0)56        conf, index = torch.max(probs, dim=0)57    58    return {59        "prediction": labels[index],60        "confidence": float(conf),61        "all_scores": {labels[i]: float(probs[i]) for i in range(5)}62    }63 64if __name__ == "__main__":65    import uvicorn66    uvicorn.run(app, host="0.0.0.0", port=7860)67