quickfix

ui/model_loader.py

@@ -21,17 +21,18 @@ class ModelLoader:
             raise ValueError(f"ClearML configuration not found for model: {modelName}")
         
         taskID = modelConfig['clearml_task_id']
-    
-        modelObject = Model(task_id=taskID)
 
         try:
-            print(f"attemtping to fetch '{modelName}' from clearML task: {taskID}")
+            print(f"Attempting to fetch '{modelName}' from ClearML task: {taskID}")
+            
+            task = Task.get_task(task_id=taskID)
 
-            modelPath = Model.get_local_copy(file_name=MODEL_ARTIFACT_NAME)
-            print(f"Weights downloaded to local path: {modelPath}")
+            modelObject = task.models["output"][0]
 
-            model = modelOne(noOfClasses=39)
+            modelPath = modelObject.get_local_copy()
+            print(f"Weights downloaded to: {modelPath}")
 
+            model = modelOne(noOfClasses=39)
             stateDict = torch.load(modelPath, map_location=self.device)
             model.load_state_dict(stateDict)
 
@@ -44,7 +45,7 @@ class ModelLoader:
             print(f"Error loading from ClearML for {modelName}: {e}")
             raise RuntimeError(f"Failed to load model from ClearML: {e}")
     
-    def loadModel(self, modelName) :
+    def loadModel(self, modelName):
         if modelName in self.modelCache:
             return self.modelCache[modelName]