Implementing Spatial Transformer Network (STN) in TensorFlow | by Parth Rajesh Dedhia | Towards Data Science

Convolution Neural Networks apply a convolution filter, on the input image in the first layer, and then on the feature maps. The CNN’s have provided extra-ordinary results by using the same weights of the filter over several parts of the same image. The results were further boosted by adding a Max Pooling layer between the Convolution layers. This pooling layer not only reduced the parameters but also improved the performance of the CNN models.
This combination of CNN and the Pooling layer provides translational in-variance: can predict a particular object only if the object is moved around in the image. However, the orientation and shape of the object should not change much. This becomes a major drawback of CNN when used in a real-world setting.
The first idea that comes to my mind is to find some black-box that could transform the image to an ideal standard so that the same classification could be re-used. Deep Mind did the same thing, they proposed a module called Spatial Transformer Network, which performs transformation to an input image. This model can be regarded as an attention module to any spatial input. Let’s have a look at the nuts and bolts of this algorithm.
This blog post is structured as follows:
Essential NumPy Operations
Image Transformations
Bi-linear interpolation
Model Design
Results and Visualization
Conclusion
The background of NumPy advance indexing makes it easier to understand Bi-Linear interpolation. If you are aware of this concept, feel free to skip this section.
>>> x = np.random.randint(10, size=8)>>> x[[1,2,5]]array([5, 6, 2])
Internal working of the above operation:
Find the value of x[1], x[2], and x[5]
Arrange them in the order in which numbers 1, 2, and 5 appear in the input list.
Return the results
This may seem very easy, when taking about a rank-1 array. Let’s see an example of a rank-2 matrix.
>>> xarray([[ 0,  1,  2,  3,  4,  5],       [ 6,  7,  8,  9, 10, 11],       [12, 13, 14, 15, 16, 17],       [18, 19, 20, 21, 22, 23]])>>> x[2,3]15
Till now stuff was pretty common, let’s send array rather then two numbers for indexing.
>>> x[[1,2,3], [3,2,3]]array([ 9, 14, 21])
Internal working of the above operation:
Finds values of x[1,3], x[2,2], and x[3,3]
Arrange them at the index location of the input list.
Return the results
Note: If the two lists are of different lengths, it will throw an IndexError.
However, in the above example we saw that the 2D array sent a 1D output. Let’s take this a step further.
>>> a = [[1,2], [3,0]]>>> b = [[0,1], [2,3]]>>> x[a, b]array([[ 6, 13],       [20,  3]])
We will discuss the internal working for one of the element of the output, rest can be derived.
Element 2 in array a is at location [0, 1]
Element 1 in the array b is at the location [0, 1]
The element in the output array is x[2, 1] and is located at [0, 1]
Some conclusion based on the above analysis:
The shape of the output array is decided by the shape of the input indexed array.
If two arrays are sent for indexing then they should always have the same shape.
That’s it for understanding the NumPy indexing. Though we will use this later in the post, understanding it here will help to focus on the idea behind Spatial Transformer Network.
Image transformations when applied to an image will change the image and produce an output. There are many transformation operations like translation, rotation, scaling, shear, and a combination of all of the above-affine transformation. Some other transformations include projective and plate-spline transformation.
Definitions
Input Image: Image on which transformation is to be applied.
Output Image: Result of transformation on an input image.
Pixel index/indices: The location of the pixel in the input image.
The matrix defined adjacent to the image is applied to the input image to produce the transformed output. I won’t be explaining the derivation as to why these matrices produce such transformation, but you can find it here.
What we need to remember is that we need a set of six numbers, which are applied to the input image, and will produce a transformed output image. But we need to understand how these transformations are applied so we will take an input image of size 300 x 300. We will apply, scale transformation in the steps below, but the same method is used for all other transformation.
Note: Transformation happens at the pixel level. We map each pixel in the output image to a pixel of the input image and copy the pixel values.
We want to find out which pixel from the input image shall be mapped to the output image at [100, 100]
We create an array x = np.asarray([100, 100, 1]) . Note, that an extra dimension is added to the input image.
Transformation matrix theta = np.asarray([[0.5, 0, 0], [0, 0.5, 0]])
Now apply matmul operation: theta @ x and get the value as array([50., 50.])
This means that the pixel value at [50, 50] in the input image will be copied to [100, 100] in the output image.
Note: The extra dimension is added to convert the coordinates from Cartesian to Homogeneous coordinates. If you are unaware of the theory, then just add the extra dimension and move ahead.
From the above maths, we can apply any affine transformation. But on contemplating the calculation done above, one would observe that we sometimes get floating-point point numbers, and pixel indices are always integers. To solve this, we use bi-linear interpolation described in the next section.
We have previously observed that the mapping of a pixel from an output image to the input image produces floating-point values, and this makes it difficult to find the pixel value at the output location. Methods like bi-cubic, bilinear, and nearest-neighbor interpolation methods are used, but we will restrict our discussion to bi-linear interpolation.
In the above image, P (a floating-point index) is the pixel index obtained after the matmul operation is applied to the output image pixel index. The set of Qxγ’s represents the pixel intensities at nearest integral pixel indices(set of all the four [x,y]). The pixel intensity value at P is calculated by the weighted average of distances from the nearest Q’s. The R1 and R2 are an intermediate representation and will help with the maths. Also, note that x2 — x1 and y2 — y1 is 1 as the distance between neighboring pixels is 1.
The denominator in all of the above equations will be 1. So the final pixel intensities at P would be given by:
The interpolated value of intensity at point P will be mapped to the outer image pixel index. We will use this formula in the NumPy implementation below.
To find the pixel intensity in the output image, one approach is to apply a for loop for each pixel index and apply the transformation and interpolate the pixel intensity at that location. However, this will be an in-efficient approach to calculate values, we will instead use the meshgrid and linspace function of NumPy. If you are unsure what they do, print out the values at intermediate steps.
The homogeneous_co_ordinates is the set of all the output image pixel indices converted to homogeneous format (ones are added to the vector). Then we apply the transformation via the matmul operation and re-arrange the output. We then separate the x and y transformed_co_ordinates and scale them to the size of image height and width. Once again, this x and y obtained after transformation correspond to the input image pixel index which will be mapped to the output image. However, these x and y are floating points and cannot directly be mapped to a pixel intensity and interpolation will be applied.
In the gist above, the x, y represent the transformed coordinates (floating-point numbers). The x0 and y0 represent the pixel index closest to but less than x and y and x1 and y1 represent the pixel index closest to but greater than x and y. The valid combinations formed by x0, x1, y0, and y1 represent the four closest integer pixel locations around x and y. The pixel intensities at these values can be equated to Qxγ.
The img contains the desired image that is to be transformed, and it should be of size [height, width]. Our objective is to find the pixel value at P, and the naive approach to solve is for each of the x and y pair, we find the corresponding 4-pixel intensities and apply the weighted average formula discussed in the previous section. However, this approach is an in-efficient one.
We instead form 4 images, where one image contains all the set of pixel values which could be found at [x0, y0], and similarly for others. The NumPy’s advance indexing which we had discussed in the first section becomes handy here. Anyone index at Ia corresponds to theQ11 for all the possible sets of floating pixel values between that index and the adjacent one. Hence, while extrapolating for a particular index, we can find the Q11 from the corresponding pixel indices at Ia. Similar operations are performed for the other 3 Q’s.
Now we directly apply the formula for bi-linear interpolation we had discussed in the previous section and find the weights for all the transformed x and y coordinates in wa, wb, wc, and wd. Now, all we are left to do is the weighted average of Ia, Ib, Ic, and Id.
Note: If you are unsure how the above set of operations performed the transformation, read the code keeping one particular location in mind, and you will discover the magic behind it. However, the best way to understand would be to visualize the intermediate images Ia, Ib, Ic, and Id by running the code.
Fun Fact: If you are keen like me and like to visualize the intermediate representation, then always convert the dtype to np.uint8. Not doing that will destroy your sleep.
If you are here with me, then pat yourself!!!
Now we have one thing remaining, each input image may need a different transformation. How do we decide which transform to apply and what should be the numbers in the corresponding transformation matrix?
The next section will help you with that.
There is an infinite number of possible transformations that could be applied to the input image and get a transformed image. To solve this, we use a set of Convolution or Dense layers and obtain the transformation matrix by re-shaping the last dense layer. We will then send this transformation matrix along with the input image to a function that will perform the bi-linear interpolation. The image obtained after transformation can be then used for the desired task.
This Spatial Transformer Module introduced by the authors at Deep Mind is categorized into three modules — Localisation Net, Grid Generator, and Sampler. However, for easy implementation, I combine the last two modules into a single module called BilinearInterpolation. The Localization net takes input an image and the BilinearInterpolation takes to input the input image and the transformation matrix obtained from the Localization net. Both of these modules are implemented by sub-classing tf.keras.layers.Layer.
Many different sets of operations could be applied to the localization net, of which I have given one set of operations. You can change them in your implementation. Only the last two layers — dense and reshape, should be always used.
The BilinearInterpolation layer takes two parameters i.e., height and width of the output map. So by changing the size of this height and width, you can produce the transformed image of a smaller size. This layer not only performs the spatial transformation, but it can also reduce the size of the image, thus performing down-sampling, without losing essential data.
If you read the above TensorFlow implemented BilinearInterpolation layer with the NumPy implementation, you can see most of the implementation remains the same. The only difference is that TensorFlow does not support direct advance indexing as elegantly as NumPy does. Hence there is a function advance_indexing which will do the operation for you.
Tada !! We have build a Spatial Transformer Module. Let’s integrate it with a Network.
From the above gist we have seen how easily these two layers can be added to our network and provide us the transformation. Now, instead of applying them at the first layer, we can replace any of the max-pooling layers with the STN. However, the localization layer needs to change as the input image may change. This flexibility allows us to use the STN module on a feature map instead of an input image.
Fun Fact: The visualizations generated in the second section are generated by using this bi-linear interpolation layer by sending fixed theta and images
The model took around 12 seconds per epoch on Colab GPU to train. The training accuracy was 99.77% and the testing accuracy was 98.86% and the model was trained for 100 epochs. However, instead of using the MNIST dataset, I have used the AffNist dataset where the images are of size 40 x 40.
In the above image, the first row of each of the rows is the input image and the second row is the image obtained after the STN network transformation on that image. We can see how the number comes to the center and also get scaled and rotated in some cases. Hence, we can conclude STN performs different transformation operation, and the values of transformation matrix are derived from the input image, thus applied only to that image.
The entire code has been posted on the repository below, feel free to clone the code and try to run the visualizations:
github.com
We have build Spatial Transformer Module, a differentiable layer, which can substitute the max-pooling layer in the contemporary networks. Since it transforms the input image/features, it acts like attention like layer to the images/features. Though this module requires a lot of background from traditional computer Vision, the elegance of this mechanism can be seen from the visualizations shown above.
I have also written a post on Capsule Network, which also dwells on the drawbacks of the Max Pooling layers and generates rich features, which in-turn generate interpretable visualizations.
M. Jaderberg, K. Simonyan, A. Zisserman, K. Kavukcuoglu, Spatial Transformer Networks, CVPR, 2015