support environment variables for preprocessing as well as post-proce…

…ssing.
vanvalenlab · Nov 18, 2020 · 4300297 · 4300297
1 parent 8fb1978
commit 4300297
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 1 deletion.
diff --git a/redis_consumer/consumers/image_consumer.py b/redis_consumer/consumers/image_consumer.py
@@ -182,7 +182,11 @@ def _consume(self, redis_hash):
             # Grap appropriate model
             model_name, model_version = utils._pick_model(label)
 
-        pre_funcs = hvals.get('preprocess_function', '').split(',')
+        if settings.LABEL_DETECT_ENABLED and label is not None:
+            pre_funcs = utils._pick_preprocess(label).split(',')
+        else:
+            pre_funcs = hvals.get('preprocess_function', '').split(',')
+
         image = self.preprocess(image, pre_funcs)
 
         # Send data to the model

diff --git a/redis_consumer/settings.py b/redis_consumer/settings.py
@@ -173,6 +173,12 @@ def _strip(x):
     2: config('CYTOPLASM_MODEL', default='FluoCytoSegmentation:0', cast=str)
 }
 
+PREPROCESS_CHOICES = {
+    0: config('NUCLEAR_PREPROCESS', default='normalize', cast=str),
+    1: config('PHASE_PREPROCESS', default='normalize', cast=str),
+    2: config('CYTOPLASM_PREPROCESS', default='normalize', cast=str)
+}
+
 POSTPROCESS_CHOICES = {
     0: config('NUCLEAR_POSTPROCESS', default='deep_watershed', cast=str),
     1: config('PHASE_POSTPROCESS', default='deep_watershed', cast=str),

diff --git a/redis_consumer/utils.py b/redis_consumer/utils.py
@@ -476,6 +476,15 @@ def _pick_model(label):
     return model.split(':')
 
 
+def _pick_preprocess(label):
+    func = settings.PREPROCESS_CHOICES.get(label)
+    if func is None:
+        logger.error('Label type %s is not supported', label)
+        raise ValueError('Label type {} is not supported'.format(label))
+
+    return func
+
+
 def _pick_postprocess(label):
     func = settings.POSTPROCESS_CHOICES.get(label)
     if func is None:

diff --git a/redis_consumer/utils_test.py b/redis_consumer/utils_test.py
@@ -411,6 +411,15 @@ def test__pick_model(mocker):
         utils._pick_model(-1)
 
 
+def test__pick_preprocess(mocker):
+    mocker.patch.object(settings, 'PREPROCESS_CHOICES', {0: 'pre'})
+    res = utils._pick_preprocess(0)
+    assert res == 'pre'
+
+    with pytest.raises(ValueError):
+        utils._pick_preprocess(-1)
+
+
 def test__pick_postprocess(mocker):
     mocker.patch.object(settings, 'POSTPROCESS_CHOICES', {0: 'post'})
     res = utils._pick_postprocess(0)