PanoViT: Vision Transformer for Room Layout Estimation from a Single Panoramic Image

The goal of the PanoViT model is to predict the room layout based on 360-degree panoramic images. According to HorizonNet [12], we represent the room layout as a one-dimensional representation. Specifically, our model takes a 360-degree panoramic image as input and outputs a $C\times 1\times W$ tensor, where $C=3$ is the number of channels, and $W$ represents the width of the panoramic image. The first two channels represent the ceiling-wall and floor-wall boundaries, and the last one represents the probability of existence of the wall-wall boundary. Based on this representation, the framework of PanoViT’s network can be divided into four modules, including a backbone, a panoramic visual transformer encoder, an edge enhancement module and a layout prediction module. A panoramic image is sent to the backbone to extract multi-scale feature maps, and sent to the edge enhancement module to obtain an edge enhancement map. The panoramic vision transformer encoder takes the original image, edge enhancement maps and multi-scale feature maps as input and outputs a feature vector for the layout prediction module to estimate the one-dimensional room layout representation. The pipeline of our network is shown in Fig. 1. We will introduce each module in detail in the next section.

The backbone is a convolutional network used to extract multi-scale features from panoramic images. In this article, we use ResNet34 as the backbone, and combine intermediate feature maps from different ResNet blocks as multi-scale features to capture low-level and high-level visual patterns. At the end of each block in ResNet34, we add a strip pooling layer [6] to make multi-scale features more robust.

The vision transformer encoder takes the original image and multi-scale feature maps as input and outputs a room layout feature vector. Fig. 2 shows the overview of the vision transformer encoder. The encoder first samples patches from the original image and the multi-scale feature maps. Then different patches are projected into different patch embeddings through different embedding operations. All the patch embeddings are attached with a recurrent position embedding to predict the room layout feature vector with a vision transformer backbone.

The structure of the layout prediction module is illustrated in Fig. 3. We first reshape the room layout feature vector back to multi-scale features following the scale of the corresponding feature maps. The features with different scales are resized to a same length with the interpolation operation, and are combined together in the vertical dimension. Then we utilize the ConvSqueezeH layer introduced in [13] to compress the multi-scale features in the vertical dimension. Finally, we apply three Conv1D layers of kernel size 3, 3, and 1 respectively with BN, ReLU to output the probability $y_{w}$ of existence of the wall-wall boundary, the ceiling-wall $y_{c}$ and the floor-wall boundaries $y_{f}$.

The edge information is useful for the wall line detection, as shown in [23]. In this paper, we extract the edge information in the frequency domain. Compared to the LSD used in [23], our method is more efficient and can achieve better accuracy.

The edge enhancement module can be represented by

where $I$ is a panorama, $E$ is the edge enhanced image, $M$ is a binary mask, $\mathscr{F}$ is the Fast Fourier transformation, and $\mathscr{F}^{-1}$ is the inverse fast Fourier transformation. Fig. 5 shows the pipeline of the edge enhancement in the frequency domain. We first compute the frequency transformation of a panoramic image $P$ with the fast Fourier transformation $\mathscr{F}$. Then we sample a part of the frequency components with a mask $M$. Finally, we reconstruct the edge enhance map $E$ with the inverse Fourier transformation $\mathscr{F}^{-1}$.

The key to edge enhancement is sampling different parts of the frequency transformation with the mask $M$. As is known, the high-frequency part contains the edge information while the low-frequency contains the mean content of the image. Therefore, a straightforward design for the mask $M$ is a high pass filter, and its reconstruction result is shown in Fig. 6 (a). As shown in [23], separately extracting the edge along different axes (i.e. x,y,z) helps to improve the accuracy. To extract the edge along different axes in the frequency domain, we propose to design the mask following the direction of the trigonometric function in the frequency transformation. The direction of the trigonometric function in the frequency transformation at $(u_{i},v_{i})$ can be calculated by

where $(u_{0},v_{0})$ is the center of the frequency transformation. Trigonometric functions in different directions highlight the edges of the panorama in the same direction. Therefore, in order to obtain edges in different directions, we design two masks for edge enhancement in the horizontal and vertical directions. The calculation formulas are respectively

and

where $\sqrt{u_{i}^{2}+v_{i}^{2}}>\theta$ works as a high pass filter, and both $\alpha$ and $\beta$ are hyperparameters. $M_{v}$ is the mask to extract the horizon edge enhanced image $E_{v}$, so we only maintain the trigonometric function whose direction angle is less than $\alpha$. The detailed information is shown in Fig. 6 (b). Finally, we concatenate two edge enhancement maps $E_{v}$ and $E_{h}$ with the original panoramic image $P$ in the channel dimension.

The loss function consists of $\mathscr{L}_{cor}$ and $\mathscr{L}_{bon}$, where $\mathscr{L}_{cor}$ measures the error between the predicted position of the corners and the ground truths. We use the binary cross-entropy to define $\mathscr{L}_{cor}$.

$\mathscr{L}_{bon}$ can be directly defined upon the errors of our two predicted ceiling and floor boundaries (i.e., $y_{f}$ and $y_{c}$) with respect to the ground truth. Previous works [12, 13] always define this error in the equirectangular image space, which does not consider the geometry information in the panorama image. In this paper, we design a panoramic geometry-aware loss in 3D space (i.e. 3D loss) inspired by RenderLoss in L2DNet [15]. In particular, we represent a pixel q on the equirectangular image with its longitude $\delta$ and latitude $\gamma$, i.e., q = ($\delta$, $\gamma$). Then we convert the pixel q to a $3$D point $P(p_{x},p_{y},p_{z})$ in the $3$D space by

where $\mathcal{F}$ denotes the convert function, $\delta$ $\in$ [-$\pi$,$\pi$], $\gamma$ $\in$ [-$\pi$/2,$\pi$/2], and $s$ denotes the scale which can be calculated by the predefined camera height. The 3D loss can be expressed as the $L_{1}$ error between the 3D point $P$ sampled from the boundary (i.e., $y_{f}$ and $y_{c}$) and the corresponding 3D point $\hat{P}$ sampled from the ground truth. An illustration of 3D loss is shown in Fig. 7.

We conduct the experiments on two different datasets, including the PanoContext and the Matterport3D (Mp3D) dataset. PanoContext dataset contains around 500 panoramas along with ground truth layout annotations. The Matterport3D dataset contains $2,295$ panoramas of the complex cases with different numbers of layout corners. We adopt the official train/val/test split in [24].

Our model is trained end-to-end with the Adam [7] optimizer, where the learning rate is $0.0001$. The hyperparameters $\alpha$, $\beta$ and $\theta$ are set to $20$, $25$ and $100$. The batch size is set to $8$ and the epoch is set to $300$. During the training stage, we apply the same data augmentations following [23], including flip, rotation, gamma transformation, and stretch. In addition, we design a new data augmentation method, adding $50$ random masks to the panoramic image, each of which is a $50\times 50$ rectangle with a value of $0$.