| import os |
| import torch |
| import torch.nn as nn |
| from torchvision.models import resnet50, ResNet50_Weights |
|
|
| class GreeneryClassifier(nn.Module): |
| """ |
| A ResNet-50 based classifier optimized for land-cover categorization. |
| Wraps the torchvision ResNet-50 architecture and replaces the final |
| fully connected layer to match the EuroSAT class count. |
| """ |
| def __init__(self, num_classes=10, pretrained=True): |
| super(GreeneryClassifier, self).__init__() |
| |
| |
| weights = ResNet50_Weights.DEFAULT if pretrained else None |
| self.model = resnet50(weights=weights) |
| |
| |
| num_ftrs = self.model.fc.in_features |
| self.model.fc = nn.Linear(num_ftrs, num_classes) |
|
|
| def forward(self, x): |
| """ |
| Forward pass for the model. |
| Args: |
| x (torch.Tensor): Input tensor of shape (Batch, 3, 64, 64). |
| Returns: |
| torch.Tensor: Logits for the 10 land-cover classes. |
| """ |
| return self.model(x) |
|
|
| def load_model(weights_path=None, num_classes=10, device=None): |
| """ |
| Utility function to initialize and optionally load weights into the classifier. |
| Args: |
| weights_path (str): Path to the .pth state dictionary. |
| num_classes (int): Number of output classes. |
| device (str, optional): Device to load the model onto ('cuda' or 'cpu'). |
| Defaults to CUDA if available, else CPU. |
| Returns: |
| GreeneryClassifier: The initialized and loaded model. |
| """ |
| if device is None: |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model = GreeneryClassifier(num_classes=num_classes) |
| if weights_path and os.path.exists(weights_path): |
| model.load_state_dict(torch.load(weights_path, map_location=device)) |
| model.to(device) |
| return model |
|
|
