Spaces:
Runtime error
Runtime error
| """ | |
| 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 | |