Update model_loader.py

remove mock implementation, local and hf model loading

ui/model_loader.py

@@ -1,207 +1,47 @@
-"""
-Model loading utilities
-Handles loading models from different sources: local files, HuggingFace, ClearML
-"""
-
 import torch
 import sys
 from pathlib import Path
+import config
 
-# Add parent directory to path to import from models
 sys.path.append(str(Path(__file__).parent.parent))
 
-from models.mock_model import MockPlantDiseaseModel, create_mock_predictions
-import config
-
 
 class ModelLoader:
-    """
-    Handles loading and managing plant disease models
-    """
-
-    def __init__(self, use_mock=True):
-        """
-        Initialize model loader
-
-        Args:
-            use_mock: If True, use mock model for development
-        """
-        self.use_mock = use_mock
-        self.model = None
+    def __init__(self):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.modelCache = {}
 
-    def load_model(self, model_name="CNN from Scratch", model_path=None):
-        """
-        Load a model based on configuration
-
-        Args:
-            model_name: Name of the model configuration
-            model_path: Optional path to model weights
-
-        Returns:
-            Loaded model
-        """
-        if self.use_mock:
-            print("Loading mock model for development...")
-            self.model = self._load_mock_model()
-        else:
-            print(f"Loading real model: {model_name}")
-            self.model = self._load_real_model(model_name, model_path)
-
-        self.model.to(self.device)
-        self.model.eval()
-        return self.model
-
-    def _load_mock_model(self):
-        """Load the mock model"""
-        model = MockPlantDiseaseModel(num_classes=len(config.CLASS_NAMES))
-        return model
-
-    def _load_real_model(self, model_name, model_path=None):
-        """
-        Load a real trained model
-
-        Args:
-            model_name: Model configuration name
-            model_path: Path to model weights
-
-        Returns:
-            Loaded model
-        """
-        model_config = config.MODEL_CONFIGS.get(model_name)
-
-        if model_config is None:
-            raise ValueError(f"Unknown model: {model_name}")
-
-        # TODO: Replace this with your actual model architecture
-        # For now, using mock model structure
-        if model_config["model_type"] == "cnn":
-            model = MockPlantDiseaseModel(num_classes=len(config.CLASS_NAMES))
-        elif model_config["model_type"] == "resnet18":
-            # TODO: Load ResNet18 transfer learning model
-            import torchvision.models as models
-            model = models.resnet18(pretrained=False)
-            model.fc = torch.nn.Linear(model.fc.in_features, len(config.CLASS_NAMES))
-        else:
-            raise ValueError(f"Unknown model type: {model_config['model_type']}")
-
-        # Load weights if path provided
-        if model_path:
-            print(f"Loading weights from {model_path}")
-            model.load_state_dict(torch.load(model_path, map_location=self.device))
-
-        return model
-
-    def load_from_clearml(self, task_id=None, project_name=None, task_name=None):
-        """
-        Load model from ClearML
+    def loadFromClearml(self, modelName):
+        modelConfig = config.MODEL_CONFIGS.get(modelName)
+        
+        if not modelConfig or 'clearml_task_id' not in modelConfig:
+            raise ValueError(f"ClearML configuration not found for model: {modelName}")
 
-        Args:
-            task_id: ClearML task ID (if known)
-            project_name: ClearML project name
-            task_name: ClearML task name
+        taskID = modelConfig['clearmml_task_id']
+        modelType  = modelConfig['modelType']
 
-        Returns:
-            Loaded model
-        """
         try:
-            from clearml import Task, Model
+            print(f"attemtping to fetch '{modelName}' from clearML task: {taskID}")
 
-            if task_id:
-                task = Task.get_task(task_id=task_id)
-            elif project_name and task_name:
-                # Get the latest task with this name
-                task = Task.get_task(
-                    project_name=project_name,
-                    task_name=task_name
-                )
-            else:
-                raise ValueError("Must provide either task_id or (project_name and task_name)")
+            modelObject = Model(taskID=taskID)
+            modelPath = modelObject.get_local_copy()
 
-            # Get the model from the task
-            model_id = task.models['output'][-1].id if task.models.get('output') else None
+            model = self.loadRealModel(modelName, modelPath, modelType)
 
-            if model_id:
-                model_obj = Model(model_id)
-                model_path = model_obj.get_local_copy()
-
-                # Load the model
-                self.model = self._load_real_model("CNN from Scratch", model_path)
-                print(f"Model loaded from ClearML task: {task_id or task_name}")
-
-                return self.model
-            else:
-                raise ValueError("No output model found in ClearML task")
-
-        except ImportError:
-            print("ClearML not installed. Install with: pip install clearml")
-            print("Falling back to mock model")
-            return self._load_mock_model()
+            return model
+        
         except Exception as e:
-            print(f"Error loading from ClearML: {e}")
-            print("Falling back to mock model")
-            return self._load_mock_model()
-
-    def load_from_huggingface(self, model_id):
-        """
-        Load model from HuggingFace Hub
-
-        Args:
-            model_id: HuggingFace model ID (e.g., "username/model-name")
-
-        Returns:
-            Loaded model
-        """
+            print(f"Error loading from ClearML for {modelName}: {e}")
+            raise RuntimeError(f"Failed to load model from ClearML: {e}")
+    
+    def loadModel(self, modelName) :
+        if modelName in self.modelCache:
+            return self.modelCache[modelName]
+        
         try:
-            from huggingface_hub import hf_hub_download
-
-            # Download model file
-            model_path = hf_hub_download(repo_id=model_id, filename="model.pth")
-
-            # Load the model
-            self.model = self._load_real_model("CNN from Scratch", model_path)
-            print(f"Model loaded from HuggingFace: {model_id}")
-
-            return self.model
-
-        except ImportError:
-            print("huggingface_hub not installed. Install with: pip install huggingface_hub")
-            print("Falling back to mock model")
-            return self._load_mock_model()
+            model = self.loadFromClearml(modelName)
+            self.modelCache[modelName] = model
+            return model
+        
         except Exception as e:
-            print(f"Error loading from HuggingFace: {e}")
-            print("Falling back to mock model")
-            return self._load_mock_model()
-
-
-def get_model(use_mock=True, **kwargs):
-    """
-    Convenience function to get a loaded model
-
-    Args:
-        use_mock: Whether to use mock model
-        **kwargs: Additional arguments for model loading
-
-    Returns:
-        Loaded model and model loader instance
-    """
-    loader = ModelLoader(use_mock=use_mock)
-    model = loader.load_model(**kwargs)
-    return model, loader
-
-
-if __name__ == "__main__":
-    # Test model loading
-    print("Testing model loading...")
-
-    # Test mock model
-    print("\n1. Loading mock model:")
-    model, loader = get_model(use_mock=True)
-    print(f"Model type: {type(model).__name__}")
-    print(f"Device: {loader.device}")
-
-    # Test with dummy input
-    dummy_input = torch.randn(1, 3, 256, 256).to(loader.device)
-    with torch.no_grad():
-        output = model(dummy_input)
-    print(f"Output shape: {output.shape}")
+            raise RuntimeError(f"Could not load model {modelName}. Check ClearML connection.")