update models init

HelmholtzAI-Consultants-Munich · Mar 7, 2024 · 03273ad · 03273ad
1 parent 726a9af
commit 03273ad
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 6 deletions.
diff --git a/src/server/dcp_server/models/custom_cellpose.py b/src/server/dcp_server/models/custom_cellpose.py
@@ -40,7 +40,7 @@ def __init__(self,
 
         # Initialize the cellpose model
         # super().__init__(**model_config["segmentor"])
-        #nn.Module.__init__(self)
+        Model.__init__(self, model_name, model_config, data_config, train_config, eval_config)
         models.CellposeModel.__init__(self, **model_config["segmentor"])
         self.model_config = model_config
         self.data_config = data_config

diff --git a/src/server/dcp_server/models/inst_to_multi_seg.py b/src/server/dcp_server/models/inst_to_multi_seg.py
@@ -49,7 +49,8 @@ def __init__(self,
         :param eval_config: Evaluation configuration.
         :type eval_config: dict
     """
-        super().__init__()
+        #super().__init__()
+        Model.__init__(self, model_name, model_config, data_config, train_config, eval_config)
 
         self.model_name = model_name
         self.model_config = model_config

diff --git a/src/server/dcp_server/models/multicellpose.py b/src/server/dcp_server/models/multicellpose.py
@@ -29,7 +29,8 @@ def __init__(self,
         :param eval_config: Evaluation configuration.
         :type eval_config: dict
         """
-
+        Model.__init__(self, model_name, model_config, data_config, train_config, eval_config)
+
         self.model_config = model_config
         self.data_config = data_config
         self.train_config = train_config
@@ -71,7 +72,6 @@ def train(self,
                     mask_class[1]!=(i+1)
                 ] = 0
                 masks_class.append(mask_class)
-            print(masks_class[0].shape)
             self.cellpose_models[i].train(imgs, masks_class)
 
         self.metric = np.mean([self.cellpose_models[i].metric for i in range(self.num_of_channels)])

diff --git a/src/server/dcp_server/models/unet.py b/src/server/dcp_server/models/unet.py
@@ -86,8 +86,9 @@ def __init__(self,
         :param eval_config: Evaluation configuration.
         :type eval_config: dict
         """
-
-        super().__init__()
+        Model.__init__(self, model_name, model_config, data_config, train_config, eval_config)
+        nn.Module.__init__(self)
+        #super().__init__()
 
         self.model_name = model_name
         self.model_config = model_config