| """ | |
| EDEN Plant Disease Detection - Mob-Res Model | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from torchvision import models | |
| import json | |
| from pathlib import Path | |
| class MobResModel(nn.Module): | |
| """ | |
| Mob-Res: MobileNetV2 + ResNet Hybrid for Disease Detection | |
| """ | |
| def __init__(self, num_classes=14, dropout=0.3): | |
| super(MobResModel, self).__init__() | |
| mobilenet = models.mobilenet_v2(weights=None) | |
| self.features = mobilenet.features | |
| in_features = 1280 | |
| self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) | |
| self.fc1 = nn.Linear(in_features, 512) | |
| self.bn1 = nn.BatchNorm1d(512) | |
| self.relu1 = nn.ReLU(inplace=True) | |
| self.dropout1 = nn.Dropout(dropout) | |
| self.fc2 = nn.Linear(512, 256) | |
| self.bn2 = nn.BatchNorm1d(256) | |
| self.relu2 = nn.ReLU(inplace=True) | |
| self.dropout2 = nn.Dropout(dropout * 0.5) | |
| self.residual_adapter = nn.Linear(in_features, 256) | |
| self.fc3 = nn.Linear(256, num_classes) | |
| def forward(self, x): | |
| x = self.features(x) | |
| x = self.avgpool(x) | |
| x = torch.flatten(x, 1) | |
| identity = x | |
| x = self.fc1(x) | |
| x = self.bn1(x) | |
| x = self.relu1(x) | |
| x = self.dropout1(x) | |
| x = self.fc2(x) | |
| x = self.bn2(x) | |
| identity = self.residual_adapter(identity) | |
| x = x + identity | |
| x = self.relu2(x) | |
| x = self.dropout2(x) | |
| x = self.fc3(x) | |
| return x | |
| def load_model(model_dir=".", device="cpu"): | |
| """Load the trained model""" | |
| with open(Path(model_dir) / "labels.json", 'r') as f: | |
| labels = json.load(f) | |
| num_classes = labels['num_classes'] | |
| model = MobResModel(num_classes=num_classes) | |
| state_dict = torch.load(Path(model_dir) / "pytorch_model.bin", | |
| map_location=device, weights_only=True) | |
| model.load_state_dict(state_dict) | |
| model.to(device) | |
| model.eval() | |
| return model, labels | |