| """ |
| model.py |
| -------- |
| EfficientNetB3-based leaf disease classifier. |
| - Pretrained on ImageNet for strong generalization |
| - Custom classification head |
| - Supports feature extraction + fine-tuning phases |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import timm |
|
|
|
|
| class LeafDiseaseModel(nn.Module): |
| """ |
| EfficientNetB3 with custom classification head for leaf disease detection. |
| |
| Architecture: |
| EfficientNetB3 backbone (ImageNet pretrained) |
| → Global Average Pooling |
| → BatchNorm → Dense(512) → GELU → Dropout(0.4) |
| → BatchNorm → Dense(256) → GELU → Dropout(0.3) |
| → Dense(num_classes) → Softmax |
| """ |
|
|
| def __init__(self, num_classes: int, pretrained: bool = True, dropout: float = 0.4): |
| super().__init__() |
| self.num_classes = num_classes |
|
|
| |
| self.backbone = timm.create_model( |
| "efficientnet_b3", |
| pretrained=pretrained, |
| num_classes=0, |
| global_pool="avg", |
| ) |
| backbone_out = self.backbone.num_features |
|
|
| |
| self.head = nn.Sequential( |
| nn.BatchNorm1d(backbone_out), |
| nn.Linear(backbone_out, 512), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.BatchNorm1d(512), |
| nn.Linear(512, 256), |
| nn.GELU(), |
| nn.Dropout(dropout * 0.75), |
| nn.Linear(256, num_classes), |
| ) |
|
|
| |
| for m in self.head.modules(): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
|
|
| |
|
|
| def freeze_backbone(self): |
| """Freeze backbone — only train the head (Phase 1).""" |
| for p in self.backbone.parameters(): |
| p.requires_grad = False |
| print(" Backbone frozen — training head only.") |
|
|
| def unfreeze_backbone(self, unfreeze_layers: int = 30): |
| """ |
| Unfreeze the last N backbone layers for fine-tuning (Phase 2). |
| EfficientNetB3 has ~360 parameters groups; last 30 covers blocks 5-7. |
| """ |
| all_params = list(self.backbone.parameters()) |
| |
| for p in all_params: |
| p.requires_grad = False |
| |
| for p in all_params[-unfreeze_layers:]: |
| p.requires_grad = True |
| trainable = sum(p.numel() for p in self.backbone.parameters() if p.requires_grad) |
| print(f" Unfrozen last {unfreeze_layers} backbone param groups " |
| f"({trainable:,} params now trainable).") |
|
|
| def unfreeze_all(self): |
| """Fully unfreeze everything (Phase 3).""" |
| for p in self.parameters(): |
| p.requires_grad = True |
| print(" All layers unfrozen.") |
|
|
| |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| features = self.backbone(x) |
| logits = self.head(features) |
| return logits |
|
|
| def get_probabilities(self, x: torch.Tensor) -> torch.Tensor: |
| """Return softmax probabilities.""" |
| return torch.softmax(self.forward(x), dim=-1) |
|
|
| def predict(self, x: torch.Tensor): |
| """Return (class_idx, confidence) tuple.""" |
| probs = self.get_probabilities(x) |
| conf, idx = torch.max(probs, dim=-1) |
| return idx, conf |
|
|
| |
|
|
| def count_parameters(self): |
| total = sum(p.numel() for p in self.parameters()) |
| trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) |
| print(f" Total parameters: {total:>12,}") |
| print(f" Trainable parameters: {trainable:>12,}") |
| return total, trainable |
|
|
|
|
| def build_model(num_classes: int, pretrained: bool = True) -> LeafDiseaseModel: |
| """Factory function — builds and returns the model.""" |
| model = LeafDiseaseModel(num_classes=num_classes, pretrained=pretrained) |
| return model |
|
|
|
|
| if __name__ == "__main__": |
| m = build_model(39) |
| m.count_parameters() |
| x = torch.randn(4, 3, 300, 300) |
| out = m(x) |
| print(f" Output shape: {out.shape}") |
|
|
