Add extract_patches as a new backend op (#11)

* Add extract_patches as a new backend op

* Add more test

* Merge master

* Fix typo

* stop using K in backend

* stop using K in backend

* Fix bug in extract_patches

* PEP8-ify

keras_contrib/backend/tensorflow_backend.py

@@ -7,6 +7,7 @@ try:
 except ImportError:
     import tensorflow.contrib.ctc as ctc
 from keras import backend as K
+from keras.backend import tensorflow_backend as KTF
 import numpy as np
 import os
 import warnings
@@ -15,6 +16,7 @@ from keras.backend.tensorflow_backend import _preprocess_conv3d_input
 from keras.backend.tensorflow_backend import _preprocess_conv3d_kernel
 from keras.backend.tensorflow_backend import _preprocess_border_mode
 from keras.backend.tensorflow_backend import _postprocess_conv3d_output
+from keras.backend.tensorflow_backend import _preprocess_border_mode
 py_all = all
 
 
@@ -65,3 +67,35 @@ def deconv3d(x, kernel, output_shape, strides=(1, 1, 1),
     x = tf.nn.conv3d_transpose(x, kernel, output_shape, strides,
                                padding=padding)
     return _postprocess_conv3d_output(x, dim_ordering)
+
+
+def extract_image_patches(X, ksizes, ssizes, border_mode="same", dim_ordering="tf"):
+    '''
+    Extract the patches from an image
+    Parameters
+    ----------
+    X : The input image
+    ksizes : 2-d tuple with the kernel size
+    ssizes : 2-d tuple with the strides size
+    border_mode : 'same' or 'valid'
+    dim_ordering : 'tf' or 'th'
+    Returns
+    -------
+    The (k_w,k_h) patches extracted
+    TF ==> (batch_size,w,h,k_w,k_h,c)
+    TH ==> (batch_size,w,h,c,k_w,k_h)
+    '''
+    kernel = [1, ksizes[0], ksizes[1], 1]
+    strides = [1, ssizes[0], ssizes[1], 1]
+    padding = _preprocess_border_mode(border_mode)
+    if dim_ordering == "th":
+        X = KTF.permute_dimensions(X, (0, 2, 3, 1))
+    bs_i, w_i, h_i, ch_i = KTF.int_shape(X)
+    patches = tf.extract_image_patches(X, kernel, strides, [1, 1, 1, 1], padding)
+    # Reshaping to fit Theano
+    bs, w, h, ch = KTF.int_shape(patches)
+    patches = tf.reshape(tf.transpose(tf.reshape(patches, [bs, w, h, -1, ch_i]), [0, 1, 2, 4, 3]),
+                         [bs, w, h, ch_i, ksizes[0], ksizes[1]])
+    if dim_ordering == "tf":
+        patches = KTF.permute_dimensions(patches, [0, 1, 2, 4, 5, 3])
+    return patches

keras_contrib/backend/theano_backend.py

@@ -1,9 +1,11 @@
 import theano
 from theano import tensor as T
 from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
+from theano.sandbox.neighbours import images2neibs
 from theano.tensor.signal import pool
 from theano.tensor.nnet import conv3d2d
 from theano.printing import Print
+
 try:
     import theano.sparse as th_sparse_module
 except ImportError:
@@ -13,6 +15,7 @@ try:
 except ImportError:
     from theano.sandbox.softsign import softsign as T_softsign
 from keras import backend as K
+from keras.backend import theano_backend as KTH
 import inspect
 import numpy as np
 from keras.backend.common import _FLOATX, floatx, _EPSILON, image_dim_ordering
@@ -21,6 +24,7 @@ from keras.backend.theano_backend import _preprocess_conv3d_kernel
 from keras.backend.theano_backend import _preprocess_conv3d_filter_shape
 from keras.backend.theano_backend import _preprocess_border_mode
 from keras.backend.theano_backend import _postprocess_conv3d_output
+
 py_all = all
 
 
@@ -73,3 +77,41 @@ def deconv3d(x, kernel, output_shape, strides=(1, 1, 1),
     conv_out = _postprocess_conv3d_output(conv_out, x, border_mode,
                                           kernel_shape, strides, dim_ordering)
     return conv_out
+
+
+def extract_image_patches(X, ksizes, strides, border_mode="valid", dim_ordering="th"):
+    '''
+    Extract the patches from an image
+    Parameters
+    ----------
+    X : The input image
+    ksizes : 2-d tuple with the kernel size
+    strides : 2-d tuple with the strides size
+    border_mode : 'same' or 'valid'
+    dim_ordering : 'tf' or 'th'
+    Returns
+    -------
+    The (k_w,k_h) patches extracted
+    TF ==> (batch_size,w,h,k_w,k_h,c)
+    TH ==> (batch_size,w,h,c,k_w,k_h)
+    '''
+    patch_size = ksizes[1]
+    if border_mode == "same":
+        border_mode = "ignore_borders"
+    if dim_ordering == "tf":
+        X = KTH.permute_dimensions(X, [0, 3, 1, 2])
+    # Thanks to https://github.com/awentzonline for the help!
+    batch, c, w, h = KTH.shape(X)
+    xs = KTH.shape(X)
+    num_rows = 1 + (xs[-2] - patch_size) // strides[1]
+    num_cols = 1 + (xs[-1] - patch_size) // strides[1]
+    num_channels = xs[-3]
+    patches = images2neibs(X, ksizes, strides, border_mode)
+    # Theano is sorting by channel
+    patches = KTH.reshape(patches, (batch, num_channels, KTH.shape(patches)[0] // num_channels, patch_size, patch_size))
+    patches = KTH.permute_dimensions(patches, (0, 2, 1, 3, 4))
+    # arrange in a 2d-grid (rows, cols, channels, px, py)
+    patches = KTH.reshape(patches, (batch, num_rows, num_cols, num_channels, patch_size, patch_size))
+    if dim_ordering == "tf":
+        patches = KTH.permute_dimensions(patches, [0, 1, 2, 4, 5, 3])
+    return patches

tests/keras_contrib/backend/test_backends.py

@@ -7,8 +7,8 @@ from keras import backend as K
 from keras.backend import theano_backend as KTH, floatx, set_floatx, variable
 from keras.backend import tensorflow_backend as KTF
 from keras_contrib import backend as KC
-from keras_contrib.backend import theano_backend as KCTH
-from keras_contrib.backend import tensorflow_backend as KCTF
+import keras_contrib.backend.theano_backend as KCTH
+import keras_contrib.backend.tensorflow_backend as KCTF
 from keras.utils.np_utils import convert_kernel
 
 
@@ -72,5 +72,33 @@ def check_composed_tensor_operations(first_function_name, first_function_args,
     assert_allclose(zth, ztf, atol=1e-05)
 
 
+class TestBackend(object):
+    def test_extract(self):
+        for input_shape in [(1, 3, 40, 40), (1, 3, 10, 10)]:
+            for kernel_shape in [2, 5]:
+                xval = np.random.random(input_shape)
+                kernel = [kernel_shape, kernel_shape]
+                strides = [kernel_shape, kernel_shape]
+                xth = KTH.variable(xval)
+                xtf = KTF.variable(xval)
+                ztf = KTF.eval(KCTF.extract_image_patches(xtf, kernel, strides, dim_ordering='th', border_mode="valid"))
+                zth = KTH.eval(KCTH.extract_image_patches(xth, kernel, strides, dim_ordering='th', border_mode="valid"))
+                assert zth.shape == ztf.shape
+                assert_allclose(zth, ztf, atol=1e-02)
+
+        for input_shape in [(1, 40, 40, 3), (1, 10, 10, 3)]:
+            for kernel_shape in [2, 5]:
+                xval = np.random.random(input_shape)
+
+                kernel = [kernel_shape, kernel_shape]
+                strides = [kernel_shape, kernel_shape]
+                xth = KTH.variable(xval)
+                xtf = KTF.variable(xval)
+                ztf = KTF.eval(KCTF.extract_image_patches(xtf, kernel, strides, dim_ordering='tf', border_mode="same"))
+                zth = KTH.eval(KCTH.extract_image_patches(xth, kernel, strides, dim_ordering='tf', border_mode="same"))
+                assert zth.shape == ztf.shape
+                assert_allclose(zth, ztf, atol=1e-02)
+
+
 if __name__ == '__main__':
     pytest.main([__file__])