Update SGN model

Author: constantinpape

development/check_uploaded_models.py

@@ -18,17 +18,23 @@
 }
 
 
-def check_segmentation_model(model_name, checkpoint_path=None):
+def check_segmentation_model(model_name, checkpoint_path=None, input_path=None, check_prediction=False, **kwargs):
     output_folder = f"result_{model_name}"
     os.makedirs(output_folder, exist_ok=True)
-    input_path = os.path.join(output_folder, f"{model_name}.tif")
-    if not os.path.exists(input_path):
-        data_path = _sample_registry().fetch(data_dict[model_name])
-        copyfile(data_path, input_path)
+    if input_path is None:
+        input_path = os.path.join(output_folder, f"{model_name}.tif")
+        if not os.path.exists(input_path):
+            data_path = _sample_registry().fetch(data_dict[model_name])
+            copyfile(data_path, input_path)
 
     output_path = os.path.join(output_folder, "segmentation.zarr")
     if not os.path.exists(output_path):
-        cmd = ["flamingo_tools.run_segmentation", "-i", input_path, "-o", output_folder, "-m", model_name]
+        cmd = [
+            "flamingo_tools.run_segmentation", "-i", input_path, "-o", output_folder, "-m", model_name,
+            "--disable_masking", "--min_size", "5",
+        ]
+        for name, val in kwargs.items():
+            cmd.extend([f"--{name}", str(val)])
         if checkpoint_path is not None:
             cmd.extend(["-c", checkpoint_path])
         subprocess.run(cmd)
@@ -38,6 +44,9 @@ def check_segmentation_model(model_name, checkpoint_path=None):
         image = imageio.imread(input_path)
         v = napari.Viewer()
         v.add_image(image)
+        if check_prediction:
+            prediction = zarr.open(os.path.join(output_folder, "predictions.zarr"))["prediction"][:]
+            v.add_image(prediction)
         v.add_labels(segmentation, name=f"{model_name}-segmentation")
         napari.run()
 
@@ -77,11 +86,13 @@ def main():
     # - Prediction works well on the GPU.
     # check_segmentation_model("IHC")
 
-    # TODO: Update model.
     # SGN segmentation (lowres):
-    # - Prediction does not work well on the CPU.
-    # - Prediction does not work well on the GPU.
-    check_segmentation_model("SGN-lowres", checkpoint_path="SGN-lowres.pt")
+    # - Prediction works well on the CPU.
+    # - Prediction works well on the GPU.
+    check_segmentation_model(
+        "SGN-lowres",
+        # boundary_distance_threshold=0.5, center_distance_threshold=None,
+    )
 
     # IHC segmentation (lowres):
     # - Prediction works well on the CPU.

development/export_models.py

@@ -21,8 +21,8 @@ def export_synapses():
 
 
 def export_sgn_lowres():
-    path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/SGN/cochlea_distance_unet_sgn-low-res-v4"  # noqa
-    model = load_model(path, device="cpu")
+    path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/SGN/cochlea_distance_unet_sgn-low-res-v5"  # noqa
+    model = load_model(path, device="cpu", name="latest")
     torch.save(model, "SGN-lowres.pt")
 
 

flamingo_tools/model_utils.py

@@ -64,14 +64,14 @@ def get_model_registry() -> None:
         "SGN": "3058690b49015d6210a8e8414eb341c34189fee660b8fac438f1fdc41bdfff98",
         "IHC": "752dab7995b076ec4b8526b0539d1b33ade5de9251aaf6863d9bd8cc9cd036b6",
         "Synapses": "2a42712b056f082b4794f15cf41b15678aab0bec1acc922ff9f0dc76abe6747e",
-        "SGN-lowres": "6accba4b4c65158fccf25623dcd0fb3b14203305d033a0d443a307114ec5dd8c",
+        "SGN-lowres": "2c773792f0ef6022c7d052c452071cf7bf45dfce6b498b408ad6cd1cc3a30d35",
         "IHC-lowres": "537f1d4afc5a582771b87adeccadfa5635e1defd13636702363992188ef5bdbd",
     }
     urls = {
         "SGN": "https://owncloud.gwdg.de/index.php/s/NZ2vv7hxX1imITG/download",
         "IHC": "https://owncloud.gwdg.de/index.php/s/wB7d2MjV5LRTP06/download",
         "Synapses": "https://owncloud.gwdg.de/index.php/s/A9W5NmOeBxiyZgY/download",
-        "SGN-lowres": "https://owncloud.gwdg.de/index.php/s/8hwZjBVzkuYhHLm/download",
+        "SGN-lowres": "https://owncloud.gwdg.de/index.php/s/OS7985CKaTTBT5g/download",
         "IHC-lowres": "https://owncloud.gwdg.de/index.php/s/EhnV4brhpvFbSsy/download",
     }
     cache_dir = get_cache_dir()

flamingo_tools/segmentation/cli.py

@@ -19,6 +19,8 @@ def _parse_kwargs(extra_kwargs, **default_kwargs):
     def _convert_argval(value):
         # The values for the parsed arguments need to be in the expected input structure as provided.
         # i.e. integers and floats should be in their original types.
+        if value is None or value == "None":
+            return None
         try:
             return int(value)
         except ValueError:
@@ -47,13 +49,21 @@ def _convert_argval(value):
 
 
 def _parse_segmentation_kwargs(extra_kwargs, model_type):
-    if model_type.startswith("SGN"):
+    if model_type == "SGN":
         default_kwargs = {
             "center_distance_threshold": 0.4,
             "boundary_distance_threshold": 0.5,
             "fg_threshold": 0.5,
             "distance_smoothing": 0.0,
-            "seg_class": "sgn_low" if model_type == "SGN-lowres" else "sgn",
+            "seg_class": "sgn",
+        }
+    elif model_type == "SGN-lowres":
+        default_kwargs = {
+            "center_distance_threshold": None,
+            "boundary_distance_threshold": 0.5,
+            "fg_threshold": 0.5,
+            "distance_smoothing": 0.0,
+            "seg_class": "sgn_low",
         }
     else:
         assert model_type.startswith("IHC")