Anmol1357/Crop_Disease

1# model_handler.py2 3import torch4from torchvision import transforms, models5from PIL import Image6from collections import OrderedDict7import os8 9# --- Configuration ---10DEVICE = "cuda" if torch.cuda.is_available() else "cpu"11MODEL_PATH = "models/best_model.pth"12 13# IMPORTANT: Ensure this list exactly matches the classes your model was trained on14CLASS_NAMES = [15    'Apple__Apple_scab', 'Apple_Black_rot', 'Apple_Cedar_apple_rust', 'Apple__healthy',16    'Blueberry__healthy', 'Cherry(including_sour)Powdery_mildew', 'Cherry(including_sour)_healthy',17    'Corn_(maize)Cercospora_leaf_spot_Gray_leaf_spot', 'Corn(maize)Common_rust', 'Corn_(maize)_Northern_Leaf_Blight',18    'Corn_(maize)healthy', 'Grape_Black_rot', 'Grape_Esca(Black_Measles)', 'Grape__Leaf_blight(Isariopsis_Leaf_Spot)',19    'Grape__healthy', 'Orange_Haunglongbing(Citrus_greening)', 'Peach__Bacterial_spot', 'Peach__healthy',20    'Pepper_bell__Bacterial_spot', 'Pepper_bell_healthy', 'Potato_Early_blight', 'Potato__Late_blight',21    'Potato__healthy', 'Raspberry_healthy', 'Rice_Bacterial_leaf_blight', 'RiceBrown_spot', 'Rice_Hispa',22    'Rice_Leaf_blast', 'RiceLeaf_scald', 'RiceNarrow_brown_leaf_spot', 'RiceNeck_blast', 'Rice_Sheath_blight',23    'Rice_healthy', 'Soybean_healthy', 'Squash_Powdery_mildew', 'Strawberry_Leaf_scorch', 'Strawberry__healthy',24    'Tomato_Bacterial_spot', 'Tomato_Early_blight', 'Tomato_Late_blight', 'Tomato_Leaf_Mold',25    'Tomato_Septoria_leaf_spot', 'Tomato_Spider_mites_Two-spotted_spider_mite', 'Tomato__Target_Spot',26    'Tomato__Tomato_Yellow_Leaf_Curl_Virus', 'Tomato__Tomato_mosaic_virus', 'Tomato__healthy',27    'Wheat_brown_rust', 'Wheat_healthy', 'Wheat_septoria'28]29 30# --- Image Transformation ---31TRANSFORM = transforms.Compose([32    transforms.Resize((224, 224)),33    transforms.ToTensor(),34    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])35])36 37def load_model():38    """Loads and returns the pre-trained EfficientNet model."""39    if not os.path.exists(MODEL_PATH):40        raise FileNotFoundError(f"Model file not found at {MODEL_PATH}.")41 42    # Initialize model43    model = models.efficientnet_v2_m(weights=None)44    model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))45 46    # Load weights47    state_dict = torch.load(MODEL_PATH, map_location=DEVICE)48    new_state_dict = OrderedDict()49    for k, v in state_dict.items():50        name = k.replace("module.", "")  # Remove DP prefix if needed51        new_state_dict[name] = v52 53    model.load_state_dict(new_state_dict, strict=False)54 55    model.to(DEVICE)56    model.eval()57    return model58 59# Load model once at startup60MODEL = load_model()61 62def get_pest_prediction(image: Image.Image) -> dict:63    """64    Analyzes an image and returns a dictionary of class probabilities.65    """66    if MODEL is None:67        raise RuntimeError("Model could not be loaded.")68 69    image_tensor = TRANSFORM(image.convert("RGB")).unsqueeze(0).to(DEVICE)70 71    with torch.no_grad():72        outputs = MODEL(image_tensor)73        probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]74 75    # Dictionary of {class_name: confidence}76    confidences = {CLASS_NAMES[i]: float(probabilities[i]) for i in range(len(CLASS_NAMES))}77    return confidences78