Merge branch 'gui_radam' of https://github.com/ian-coccimiglio/cellpose…

… into ian-coccimiglio-gui_radam
MouseLand · Oct 22, 2024 · af081e7 · af081e7
2 parents 8389f9f + 0d59d07
commit af081e7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/cellpose/gui/gui.py b/cellpose/gui/gui.py
@@ -295,6 +295,7 @@ def __init__(self, image=None, logger=None):
             "learning_rate": 0.1,
             "weight_decay": 0.0001,
             "n_epochs": 100,
+            "SGD": True,
             "model_name": "CP" + d.strftime("_%Y%m%d_%H%M%S"),
         }
 
@@ -2170,13 +2171,15 @@ def train_model(self, restore=None, normalize_params=None):
         save_path = os.path.dirname(self.filename)
 
         print("GUI_INFO: name of new model: " + self.training_params["model_name"])
+        print(f"GUI_INFO: SGD activated: {self.training_params['SGD']}")
         self.new_model_path, train_losses = train.train_seg(
             self.model.net, train_data=self.train_data, train_labels=self.train_labels,
             channels=self.channels, normalize=normalize_params, min_train_masks=0,
-            save_path=save_path, nimg_per_epoch=max(8, len(self.train_data)), SGD=True,
+            save_path=save_path, nimg_per_epoch=max(8, len(self.train_data)),
             learning_rate=self.training_params["learning_rate"],
             weight_decay=self.training_params["weight_decay"],
             n_epochs=self.training_params["n_epochs"],
+            SGD=self.training_params["SGD"],
             model_name=self.training_params["model_name"])[:2]
         # save train losses
         np.save(str(self.new_model_path) + "_train_losses.npy", train_losses)

diff --git a/cellpose/gui/guiparts.py b/cellpose/gui/guiparts.py
@@ -233,6 +233,12 @@ def __init__(self, parent, model_strings):
             self.edits[-1].setFixedWidth(200)
             self.l0.addWidget(self.edits[-1], i + yoff, 1, 1, 1)
 
+        yoff += 1
+        use_SGD = "SGD"
+        self.useSGD = QCheckBox(f"{use_SGD}")
+        self.useSGD.setChecked(True)
+        self.l0.addWidget(self.useSGD, i+yoff, 1, 1, 1)
+
         yoff += len(labels)
 
         yoff += 1
@@ -289,6 +295,7 @@ def accept(self, parent):
             "weight_decay": float(self.edits[1].text()),
             "n_epochs": int(self.edits[2].text()),
             "model_name": self.edits[3].text(),
+            "SGD": True if self.useSGD.isChecked() else False,
             #"use_norm": True if self.use_norm.isChecked() else False,
         }
         self.done(1)