quickfix

ui/config.py

@@ -4,12 +4,12 @@ CLEARML_TASK_NAME_DEFAULT = "CNN Training (Latest)"
 MODEL_CONFIGS = {
     "Shallow CNN": {
         "description": "modelOne trained on 20 epochs",
-        "model_type" : "cnn",
+        "class" : "modelOne",
         "clearml_task_id": "dca82d7c2f404c249f2e5325aaf77207"
     },
     "Deep CNN": {
         "description" : "modleTwo trained on 30 epochs",
-        "model_type": "cnn",
+        "class": "betterCNN",
         "clearml_task_id": "c79a6939b46a4882a7fdaee117b1f32e"
     }
 }

ui/model_loader.py

@@ -2,12 +2,17 @@ import torch
 import sys
 from pathlib import Path
 import config
-import utils
 from clearml import Task
 from models.modelOne import modelOne
+from models.modelTwo import BetterCNN
 
 sys.path.append(str(Path(__file__).parent.parent))
 
+MODEL_CLASSES = {
+    "modelOne": modelOne,
+    "betterCNN": BetterCNN
+}
+
 MODEL_ARTIFACT_NAME = 'best_model'
 
 class ModelLoader:
@@ -22,6 +27,7 @@ class ModelLoader:
             raise ValueError(f"ClearML configuration not found for model: {modelName}")
         
         taskID = modelConfig['clearml_task_id']
+        className = modelConfig['class']
 
         try:
             print(f"Attempting to fetch '{modelName}' from ClearML task: {taskID}")
@@ -29,16 +35,13 @@ class ModelLoader:
             task = Task.get_task(task_id=taskID)
             print("Available artifacts:", task.artifacts.keys())
 
-            # Fetch the artifact 'model_one.pt'
             artifact = task.artifacts.get(MODEL_ARTIFACT_NAME)
-
             if artifact is None:
                 raise RuntimeError(
                     f"Artifact '{MODEL_ARTIFACT_NAME}' not found in ClearML task {taskID}"
                 )
 
             modelPath = artifact.get_local_copy()
-
             if modelPath is None:
                 raise RuntimeError(
                     f"Artifact '{MODEL_ARTIFACT_NAME}' could not be downloaded (returned None)"
@@ -46,8 +49,11 @@ class ModelLoader:
 
             print(f"Weights downloaded to: {modelPath}")
 
-            # Load PyTorch model
-            model = modelOne(noOfClasses=39)
+            # Load correct model class
+            ModelClass = MODEL_CLASSES[className]
+            model = ModelClass(noOfClasses=39)
+
+            # Load weights
             stateDict = torch.load(modelPath, map_location=self.device)
             model.load_state_dict(stateDict)