add model loading to modelLoader

ui/config.py

@@ -6,10 +6,5 @@ MODEL_CONFIGS = {
         "description": "Custom CNN model trained from scratch",
         "model_type": "cnn",
         "clearml_task_id": "fe14662da63d45bf9208fdf9856d2fcc" 
-    },
-    "Transfer Learning (ResNet18)": {
-        "description": "Fine-tuned ResNet18 model",
-        "model_type": "resnet18",
-        "clearml_task_id": "SET_ME_TO_YOUR_RESNET_TASK_ID"
     }
 }

ui/model_loader.py

@@ -3,6 +3,7 @@ import sys
 from pathlib import Path
 import config
 from clearml import Model
+from ..models import modelOne
 
 sys.path.append(str(Path(__file__).parent.parent))
 
@@ -26,8 +27,15 @@ class ModelLoader:
 
             modelObject = Model(model_id=taskID)
             modelPath = modelObject.get_local_copy()
+            print(f"Weights downloaded to local path: {modelPath}")
 
-            model = self.loadRealModel(modelName, modelPath, modelType)
+            model = self.modelOne()
+
+            stateDict = torch.load(modelPath, map_location=self.device)
+            modelObject.load_state_dict(stateDict)
+
+            model.to(self.device)
+            model.eval()
 
             return model