Feature/plugin updates (AllenCellModeling#345)

* add timeout * add autothreshold as default postprocessing; * update with Dannys comments * precommit --------- Co-authored-by: Benjamin Morris <[email protected]>
KartikeyKansal1 · Mar 6, 2024 · 568fdc3 · 568fdc3
1 parent 1c11e5b
commit 568fdc3
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/configs/model/im2im/segmentation_plugin.yaml b/configs/model/im2im/segmentation_plugin.yaml
@@ -27,10 +27,8 @@ task_heads:
         _target_: cyto_dl.models.im2im.utils.postprocessing.ActThreshLabel
         rescale_dtype: numpy.uint8
       prediction:
-        _target_: cyto_dl.models.im2im.utils.postprocessing.ActThreshLabel
-        activation:
-          _target_: torch.nn.Sigmoid
-        rescale_dtype: numpy.uint8
+        _target_: cyto_dl.models.im2im.utils.postprocessing.AutoThreshold
+        method: "threshold_otsu"
     save_input: True
 
 optimizer:

diff --git a/configs/trainer/default.yaml b/configs/trainer/default.yaml
@@ -17,3 +17,5 @@ check_val_every_n_epoch: 1
 # makes training slower but gives more reproducibility than just setting seeds
 deterministic: False
 detect_anomaly: False
+
+max_time: null
diff --git a/cyto_dl/models/im2im/utils/postprocessing/auto_thresh.py b/cyto_dl/models/im2im/utils/postprocessing/auto_thresh.py
@@ -1,9 +1,9 @@
 import importlib
-from typing import Union
+from typing import Optional, Union
 
 
 class AutoThreshold:
-    def __init__(self, method: Union[float, str]):
+    def __init__(self, method: Optional[Union[float, str]] = None):
         if isinstance(method, float):
 
             def thresh_func(image):
@@ -14,9 +14,13 @@ def thresh_func(image):
                 thresh_func = getattr(importlib.import_module("skimage.filters"), method)
             except AttributeError:
                 raise AttributeError(f"method {method} not found in skimage.filters")
+        elif method is None:
+            thresh_func = None
         else:
             raise TypeError("method must be a float or a string")
         self.thresh_func = thresh_func
 
     def __call__(self, image):
+        if self.thresh_func is None:
+            return image
         return image > self.thresh_func(image)