Albumentations Lambda support

Author: qubvel

segmentation_models_pytorch/encoders/__init__.py

@@ -45,7 +45,7 @@ def get_preprocessing_fn(encoder_name, pretrained='imagenet'):
     mean = settings[pretrained].get('mean')
     std = settings[pretrained].get('std')
 
-    def _preprocess_input(x):
-        return preprocess_input(x, mean=mean, std=std, input_space=input_space, input_range=input_range)
+    def _preprocess_input(x, **kwargs):
+        return preprocess_input(x, mean=mean, std=std, input_space=input_space, input_range=input_range, **kwargs)
 
     return _preprocess_input

segmentation_models_pytorch/encoders/_preprocessing.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 
-def preprocess_input(x, mean=None, std=None, input_space='RGB', input_range=None):
+def preprocess_input(x, mean=None, std=None, input_space='RGB', input_range=None, **kwargs):
 
     if input_space == 'BGR':
         x = x[..., ::-1].copy()