Spaces:
Runtime error
Runtime error
File size: 2,914 Bytes
f4bd707 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """
utils/model_loader.py
Loads all 4 trained models from the models/ directory.
Models are expected to be PyTorch .pt or .pth files.
"""
import torch
import torchvision.models as models
import os
def build_model(num_classes: int) -> torch.nn.Module:
"""
Build a ResNet-50 model with a custom final classification layer.
We use ResNet-50 as the backbone — a powerful, proven architecture.
The final fully-connected layer is replaced to match our number of classes.
Args:
num_classes: How many categories this model predicts.
Returns:
A PyTorch model (not yet loaded with weights).
"""
model = models.resnet50(weights=None) # No pretrained weights — we'll load ours
# Replace the final layer to match our class count
in_features = model.fc.in_features
model.fc = torch.nn.Linear(in_features, num_classes)
return model
def load_model(path: str, num_classes: int, device: torch.device) -> torch.nn.Module:
"""
Load a trained model from a .pt or .pth file.
Args:
path: File path to the saved model weights.
num_classes: Number of output classes for this model.
device: CPU or GPU.
Returns:
Loaded model ready for inference.
"""
if not os.path.exists(path):
raise FileNotFoundError(
f"Model file not found: '{path}'\n"
f"Make sure your trained model is saved at this location."
)
model = build_model(num_classes)
# Load the saved weights into the model
state_dict = torch.load(path, map_location=device)
model.load_state_dict(state_dict)
model.to(device)
model.eval() # Set to evaluation mode (disables dropout, batchnorm training behavior)
return model
def load_all_models() -> dict:
"""
Load all 4 models and return them in a dictionary.
Model files expected in models/ folder:
- tree_vs_nontree.pt (2 classes: Tree, Non-Tree)
- species.pt (2 classes: Mango, White Gum)
- mango_stage.pt (4 classes: Seedling, Sapling, Mature, Overmature)
- gum_stage.pt (4 classes: Seedling, Sapling, Mature, Overmature)
Returns:
Dictionary with keys: tree_vs_nontree, species, mango_stage, gum_stage
"""
# Use GPU if available, otherwise CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Define model configurations: (file_path, number_of_classes)
model_configs = {
"tree_vs_nontree": ("models/tree_vs_nontree.pt", 2),
"species": ("models/species.pt", 2),
"mango_stage": ("models/mango_stage.pt", 4),
"gum_stage": ("models/gum_stage.pt", 4),
}
loaded_models = {}
for name, (path, num_classes) in model_configs.items():
loaded_models[name] = load_model(path, num_classes, device)
return loaded_models
|