Spaces:
Sleeping
Sleeping
| import torch | |
| import logging | |
| from pathlib import Path | |
| from app.models.unet_model import BuildUNet | |
| from app.core.config import settings | |
| from app.core.exceptions import ModelLoadError | |
| logger = logging.getLogger(__name__) | |
| class ModelLoader: | |
| def __init__(self): | |
| self.model = None | |
| self.device = None | |
| self._load_device() | |
| def _load_device(self): | |
| if settings.MODEL_DEVICE == "cuda" and torch.cuda.is_available(): | |
| self.device = torch.device("cuda") | |
| logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}") | |
| else: | |
| self.device = torch.device("cpu") | |
| logger.info("Using CPU") | |
| def load_model(self): | |
| try: | |
| model_path = Path(settings.MODEL_PATH) | |
| if not model_path.exists(): | |
| raise ModelLoadError(f"Model file not found at {settings.MODEL_PATH}") | |
| logger.info(f"Loading model from {settings.MODEL_PATH}") | |
| self.model = BuildUNet(num_classes=settings.NUM_CLASSES) | |
| checkpoint = torch.load(model_path, map_location=self.device) | |
| if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: | |
| self.model.load_state_dict(checkpoint["model_state_dict"]) | |
| else: | |
| self.model.load_state_dict(checkpoint) | |
| self.model.to(self.device) | |
| self.model.eval() | |
| logger.info("Model loaded successfully") | |
| except Exception as e: | |
| logger.error(f"Failed to load model: {str(e)}") | |
| raise ModelLoadError(f"Failed to load model: {str(e)}") | |
| def get_model(self): | |
| if self.model is None: | |
| raise ModelLoadError("Model not loaded. Call load_model() first.") | |
| return self.model | |
| def get_device(self): | |
| return self.device | |
| model_loader = ModelLoader() | |