A Simple Approach to Image Tilt Correction with Self-Attention MobileNet for Smartphones

Squeeze-and-Excite [Hu et al.(2018)Hu, Shen, and Sun]: The squeeze operation captures channel-wise dependencies in a feature map by using global average pooling operation along the channel dimension to aggregate information of each 2D feature map through a scalar value that results in a channel descriptor vector. During the excite operation, we learn channel-wise dependencies by passing the channel descriptor vector through a 2-layer neural network which outputs channel-wise attention weights. This attention vector is then multiplied by the input feature map to adaptively weight each channel and improve the representational power of the feature maps.

Spatial Self-Attention: Standard convolutional kernels process a local neighborhood of an image at a time because of the smaller receptive field sizes for the input feature maps. Although, the kernels with large receptive fields can cover more region but that will incur high computational costs because of the increased number of parameters and training time. Self-attention modules can complement convolutional layers by learning global or long-range dependencies between different image regions without much computational overhead. To learn spatial self-attention, we take the feature map vectors along the channel dimension and calculate their key, query, and value representations in a low-dimensional subspace. The dot product of all key vectors with every other query vector with the softmax function gives us the attention weights between all the regions of the image. More specifically, after squeeze and excite operation, let the feature map be $\mathbf{F\in\mathbb{R}^{H\times W\times C}}$ and after expanding it along spatial dimensions we get the matrix $\mathbf{F\in\mathbb{R}^{N\times C}}$, which contains $N(=H\cdot W)$, $C$-dimensional vectors representing various image regions. We use learnable weight matrices $\mathbf{W_{K}\in\mathbb{R}^{C\times C`}},\mathbf{W_{Q}\in\mathbb{R}^{C\times C`}}$, and $\mathbf{W_{v`}\in\mathbb{R}^{C\times C`}}$ to calculate key($\mathbf{K}$), query($\mathbf{Q}$) and value($\mathbf{v`}$) vectors respectively.

Here $C`=C/r$, $r>1$, $r$ is the reduction ratio which is used to decrease the dimensions of the vectors and calculate attention weights and values in a low-dimensional subspace. Let $\mathbf{F^{N\times C}}\equiv\mathbf{F}$

	$\displaystyle\mathbf{K}=\mathbf{F\cdot W_{K}},\quad\mathbf{Q}$	$\displaystyle=\mathbf{F\cdot W_{Q}},\quad\mathbf{v`}=\mathbf{F\cdot W_{v`}}$		(1)
	$\displaystyle\mathbf{\beta^{N\times N}}=\mathbf{Q\cdot K^{T}}\implies a_{ij}$	$\displaystyle=\frac{\exp(\beta_{ij})}{\sum_{j=1}^{N}\exp(\beta_{ij})},i,j=1,\dots,N$		(2)
	$\displaystyle\mathbf{V^{N\times C`}}$	$\displaystyle=\mathbf{a}\cdot\mathbf{v`}$		(3)

Here, $\mathbf{a\in\mathbb{R}^{N\times N}}$ is the self-attention matrix, and $a_{ij}$ is the attention weight that region $i$ puts on region $j$. $\mathbf{V\in\mathbb{R}^{N\times C`}}$in Eq.3 is a value matrix in the low-dimensional subspace and we use $\mathbf{W_{V}\in\mathbb{R}^{C`\times C}}$ to project it into the original subspace to get the self-attention feature maps, $\mathbf{S\in\mathbb{R}^{N\times C}}$. Finally, we will add a residual connection to get the final output $\mathbf{\tilde{F}^{N\times C}}$.

	$\displaystyle\mathbf{S}$	$\displaystyle=\mathbf{V}\cdot\mathbf{W_{V}}$		(4)
	$\displaystyle\mathbf{\tilde{F}}$	$\displaystyle=\mathbf{F}+\alpha\mathbf{S}$		(5)

where $\alpha$ is a trainable scalar parameter initialized to $0$. Eq.5 implies that the model first learns image features around the local neighbourhood and then gradually moves on to learn global dependencies [Zhang et al.(2019)Zhang, Goodfellow, Metaxas, and Odena]. The combination of squeeze and excite operation and spatial self-attention gives us the feature maps that are rich in content as well as contextual information.

The learning of global long-range dependencies and relative positioning of different image regions improved the model predictions on image tilt detection tasks. From the heatmaps, in Fig.4, we can see that the query (red) region is able to shift its attention with the image orientation. Specifically, in Fig.4-[1.a, 2.a, 3.a] (left column), the query (red) region on the horizon was able to focus on the other horizon points despite different orientations of the same image. In Fig.4-SET-B (right column), there is no clear horizon line but the query regions are able to attend to relevant regions for tilt detection. For example, query region on the ground in heatmaps-[4.b, 5.b] is attending to other points on the ground and query region in Fig.4-[4.a, 6.b] in sky is focusing on other points above the skyline. Also note that query regions in Fig.4-[1.c, 2.c, 3.c] (left column) are able to attend to far locations which indicates learning of long-range dependencies and spatial information. Therefore, for detecting image tilt angle, the neural network model needs to learn the relative positioning of the image pixels to differentiate between various distinct image orientations.

For the MobileNetV3 model, the input image size is $224\times 224$ and the subsequent layers decrease the feature map size to $7\times 7$ in strides of $2$. We integrated Spatial Self-Attention Modules within MobileNetV3 to get Self-Attention MobileNet. Spatial Self-Attention Modules were applied to the blocks of sizes $56\times 56$, $28\times 28$, $14\times 14$, and $7\times 7$. We selected these blocks because they are high-level image features which encodes meaningful image representations. Furthermore, the computational costs of calculating self-attention on these blocks was not very high because of the small feature map sizes. We also replaced the 1280-dimensional fully connected layer of MobileNetV3 with a 720-dimensional layer. This helped in reducing the MAdds that were added due to the introduction of spatial self-attention operations. As a result, the proposed SA-MobileNet model, when converted to Tflite, came out faster by an average of 4ms than the corresponding MobileNetV3 Tflite model when tested on Snapdragon 750 Octa core. All the convolutional kernels and fully connected layers were initialized from pretrained ImageNet [Krizhevsky et al.(2012)Krizhevsky, Sutskever, and Hinton] weights apart from the newly added self-attention blocks that were initialized randomly.

The network was trained end-to-end using RMSprop [Tieleman and Hinton(2012)] optimizer with momentum $0.9$. The initial learning rate was $0.001$ and it was decayed using exponential learning rate schedule with $40k$ decay steps and $0.95$ decay rate.

We used publicly available SUN397 [Xiao et al.(2010)Xiao, Hays, Ehinger, Oliva, and Torralba], ADE20K [Zhou et al.(2017)Zhou, Zhao, Puig, Fidler, Barriuso, and Torralba], and NYU-V1 [Silberman and Fergus(2011)] datasets. SUN397 is a scene understanding dataset that contains 397 well-sampled categories of diverse scenes with $108,754$ distinct images. There are 10 train-test partitions for the dataset, and for our evaluation dataset, we took the union of images from all the test partitions that resulted in $15,691$ images. The remaining $92,793$ images were used for training. ADE20K dataset is another scene parsing dataset with $20,210$ images in the training set and $5,000$ images in the evaluation set. SUN397 and ADE20K datasets contain wide variety of natural images that may or may not contain straight vertical or horizontal lines. That is why we set the interval length, $I=2\degree$ while training on these two datasets. We also used NYU-V1 Depth dataset that contains frames from video sequences of various indoor scenes recorded from both RGB and Depth camera of Microsoft Kinect. Before using this dataset to train our model, we had to make the images upright because all the images were skewed as seen from Fig.6. We straightened the set of $2,282$ images and used $2000$ images for our training and $282$ images for testing. We split the data in such a way that frames from the same indoor scene does not come in both the splits. We set the interval length $I=1\degree$ because this dataset was manually annotated and we saw highly accurate results on the evaluation data.

We train MobileNetV3 model as a baseline for mobile devices on SUN397, ADE20K and NYU-V1 dataset using the training approach, described in section 3.2. The proposed SA-MobileNet consistently performs better in terms of detection accuracies and angle errors on all the evaluation datasets as seen from Fig.5 and Table.4. We also trained MobileNetV3 and SA-MobileNet for regressing the image tilt angle, by using angle loss function, given by Eq.7 and 8, and AdaDelta optimizer on ADE20k dataset. The SA-MobileNet model gives us lower angle error when compared to the MobileNetV3 model as shown in Fig 5.d and Table 3.

	$\displaystyle e$	$\displaystyle=\mathinner{\!\left\lvert a_{true}-a_{pred}\right\rvert}$		(7)
	$\displaystyle\mathbf{\mathcal{L}_{angle}}$	$\displaystyle=min\{e,360-e\}$		(8)

Here $a_{true}$ and $a_{pred}$ are groundtruth and predicted angles respectively, with values ranging from $0\degree-to-359\degree$. We also use this loss function to calculate the angle errors of the model trained with the multi-label approach.

From Table.4 we see that the proposed SA-MobileNet model resulted in very low-angle errors on NYU-V1, ADE20K, and SUN397 dataset when compared with MobileNetV3 model. From the evaluation accuracy plots in Fig.5.a, Fig.5.b, Fig.5.c, we can see that the accuracy curve of our model (blue) is above the curve of MobileNetV3 model (green) during training for all the datasets. In Fig.5.d, we plot angle errors for both the models which were trained for regressing the exact orientation angle on ADE20K dataset. The SA-MobileNet model produced low angle errors when compared to the MobileNetV3 model. This also justifies that the long-range dependencies learned by the self-attention modules are necessary for image tilt detection. We also trained the ResNet-50 model as a baseline to validate the effectiveness of our training pipeline. Though, the ResNet-50 model outperforms both the MobileNetV3 and SA-MobileNet, the improvement is marginal despite the huge difference in the number of parameters between the ResNet-50 and the light-weight models.

Comparison with previous works: We also present comparison of our training approach with previous works for this problem in Table 1. Over the past few years many different methods have been proposed to tackle this problem but they have high angle errors or low accuracies. ResNet-50 baseline model gives the lowest angle error on the test dataset to give state-of-the-art results. However, the ResNet-50 network cannot be deployed on low-resource devices because it has over 25 million parameters. But the proposed Self-Attention MobileNet has around 4 million parameters with low-latency values and state-of-the-art results for mobile devices that makes it ideal to be deployed in smartphones for real-time inferences.