File size: 1,907 Bytes
43abac3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | 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__()
# Load pre-trained ResNet50
weights = ResNet50_Weights.DEFAULT if pretrained else None
self.model = resnet50(weights=weights)
# Replace the final fully connected layer
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
|