Adapted preprocessing. (#60)

Author: martinsbruveris

tests/models/test_factory.py

@@ -125,11 +125,13 @@ def test_model_path(model_name):
 
 
 @pytest.mark.parametrize("model_name", TEST_ARCHITECTURES)
-@pytest.mark.parametrize("input_shape", [(8, 8, 3), (1, 4, 4, 3)])
+@pytest.mark.parametrize("input_size", [(8, 8), (1, 4, 4)])
+@pytest.mark.parametrize("in_channels", [1, 3, 5, 6])
 @pytest.mark.parametrize("dtype", ["float32", "float16"])
-def test_preprocessing(model_name, input_shape, dtype):
+def test_preprocessing(model_name, input_size, in_channels, dtype):
+    input_shape = (*input_size, in_channels)
     img = tf.ones(input_shape, dtype)
-    preprocess = create_preprocessing(model_name, dtype)
+    preprocess = create_preprocessing(model_name, in_channels=in_channels, dtype=dtype)
     img = preprocess(img)
     assert img.shape == input_shape
     assert img.dtype == dtype

tfimm/models/factory.py

@@ -111,14 +111,20 @@ def create_model(
     return model
 
 
-def create_preprocessing(model_name: str, dtype: Optional[str] = None) -> Callable:
+def create_preprocessing(
+    model_name: str,
+    *,
+    in_channels: Optional[float] = None,
+    dtype: Optional[str] = None,
+) -> Callable:
     """
     Creates a function to preprocess images for a particular model.
 
     The input to the preprocessing function is assumed to be values in range [0, 255].
 
     Args:
         model_name: Model for which to create preprocessing function.
+        in_channels: Number of input channels to model
         dtype: Output dtype.
 
     Returns:
@@ -130,9 +136,22 @@ def create_preprocessing(model_name: str, dtype: Optional[str] = None) -> Callab
     cfg = model_config(model_name)
     dtype = dtype or tf.keras.backend.floatx()
 
+    def _adapt_vector(v, n):
+        """Adapts vector v to length n by repeating as necessary."""
+        v = tf.convert_to_tensor(v, dtype=dtype)
+        m = tf.shape(v)[0]
+        nb_repeats = n // m + 1
+        v = tf.tile(v, [nb_repeats])
+        v = v[:n]
+        return v
+
+    in_channels = in_channels or cfg.in_channels
+    mean = _adapt_vector(cfg.mean, in_channels)
+    std = _adapt_vector(cfg.std, in_channels)
+
     def _preprocess(img: tf.Tensor) -> tf.Tensor:
         img = tf.cast(img, dtype=dtype) / 255.0
-        img = (img - cfg.mean) / cfg.std
+        img = (img - mean) / std
         return img
 
     return _preprocess