OAS-Net: Occlusion Aware Sampling Network for Accurate Optical Flow

Given input frames of $I_{1},I_{2}$, our OAS-Net estimates optical flow in a coarse-to-fine manner like many approaches [6, 8, 12, 10], where a feature extraction network is adopted to build pyramid features $f_{1}^{k},f_{2}^{k},k=1,2,\ldots,6$. Each pyramid scale in OAS-Net contains two convolution layers and the first stride 2 convolution down samples spatial resolutions. Commonly, we set an incremental 16, 32, 64, 96, 128 and 160 feature maps from level 1 to level 6 respectively.

As shown in Fig 2, 5 sub-networks sharing the same structure for occlusion aware optical flow estimation are employed in each level to refine $2\times$ upsampled flow fields together with occlusion awareness maps. For completeness, we first describe the shared optical flow and occlusion decoder. Since there is no dilated context network [6, 9] utilized by OAS-Net, we set a sequential connected decoder with three more convolution layers to enlarge receptive field for fair comparison, which outputs 128, 128, 128, 128, 128, 96, 64 and 32 channels. The final flow and occlusion maps are predicted by two separate convolution heads making OAS-Net lightweight and efficient. To keep meaningful range of occlusion awareness map, we add a non-linear sigmoid activation to restrict output from 0 to 1. Unlike [9] which relys on both flow and occlusion ground-truth for co-training, our model only requires optical flow label, while the occlusion awareness is automatically learned by flow supervision.

One key step in optical flow approaches is to build cost volume which can provide better correspondence representation than direct convolution feature. To deal with the challenging large displacement, mainstream solution is to leverage bilinear sampling based warping operation [7, 6, 8, 9, 12] to reduce offset between source and target corresponding feature according to the roughly estimated flow field. When the target feature is warped to the source one, followed matching cost volume can be built by calculating feature similarity within a local searching area. As shown in Fig 1c, this warping based cost volume can be formulated as:

Where $f_{warp}^{k}$ means warped target feature in level $k$, $\bf x$ means spatial coordinates and $\bf d$ represents searching offsets. $\cdot$ is to calculate innner product for similarity measurement.

As mentioned above, the warping based cost volume suffers from duplicate artifacts, which is also known as ghosting. To solve the problem, we build cost volume with a novel sampling based approach as shown in Fig 1d that avoids warping and ghosting issues. This operation can be formulated as:

Where the new vector $\bf f$ denotes previous optical flow field in the location of coordinate $\bf x$. Different from a dilated differential volume in [11], we adopt inner product based volume whose $\bf d$ traverses a searching square with radius 4. Compared to Eq 1, the flow based warping operation is replaced by directly sampling flow guided searching grids in the target pyramid feature. Experiments will show proposed sampling based correlation method helps to improve performance.

Although sampling based cost volume can bypass annoying ghosting feature, another challenging occlusion problem is unavoidable. As shown in Fig 1, the red cross pointing bush in the source image is covered by the moving arm in the target image. This type of occluded region is unrecoverable no matter to take either correlation methods.

In regard to supervised methods [4, 5, 7, 6, 8, 12, 10, 11], little attention has been paid to handle occlusion explicitly during pure supervison by densely labeled optical flow. To endow flow estimation with occlusion awareness, we present a novel occlusion aware module to better handle two types of matching regions in raw cost volume separately, which can be seamlessly embedded into the coarse-to-fine structure.

As depicted in Fig 2, there are four steps to build occlusion aware cost volume. First, previous $2\times$ upsampled occlusion awareness is subtracted from a full one tensor to create a reversed occlusion awareness map. Second, the complementary occlusion maps separately multiply raw cost volume to extract two occlusion reweighted cost volumes. Third, two disparate convolution layers are employed to filter above occlusion aware cost volumes respectively. Finally, two filtered matching costs are merged by addition following a leaky ReLU layer. The whole process can be formulated as:

Where $O(\bf x)$ means occlusion awareness map, $\otimes$ represents broadcast multiplication and $\oplus$ stands for element-wise addition. Referring to Eq 3, proposed occlusion awareness map can also be interpreted as occlusion probability map, which can be viewed as one type of self-attention mechanism [13]. We will visualize occlusion awareness and demonstrate our occlusion aware cost volume can improve optical flow accuracy, especially in fast-moving scenarios.

We adopt the two stage training schedule in FlowNet2 [5]. OAS-Net is firstly trained on FlyingChairs dataset [4] with $S_{short}$ learning schedule, i.e., initial learning rate is set to $1e-4$ and decays half at $300k,400k$ and $500k$ iterations with total $600k$ iterations. Then, the pretrained model is fine-tuned on FlyingThings3D dataset [14] following the $S_{fine}$ schedule, i.e., with initial learning rate $1e-5$ and decays half at $200k,300k$ and $400k$ iterations with total $500k$ iterations. For FlyingChairs, batch size is set to 8 and crop size is set to $320\times 448$, and for FlyingThings3D, we take batch size 4 and adopt $384\times 768$ crops. Multiple data augmentation methods, including mirror, translate, rotate, zoom, squeeze and color jitter are employed to enrich training distribution and prevent overfitting. Since pyramid structure of OAS-Net is the same with PWC-Net [6] and LiteFlowNet [8], we use the same multi-scale L2 loss [6] for supervised learning. Adam [15] optimizer is adopted in all stages, and our experiments are conducted on one NVIDIA GTX 1080Ti GPU with PyTorch.

To explore and verify value of proposed approaches, we conduct ablation with 4 variations based on pairwise combinations. For sake of effective comparison, all variants follow above two stage training schedule and are evaluated on Sintel Final, KITTI 2012 training datasets, as listed in Table 1.

Compared with the baseline method, which adopts warping based correlation and does not include occlusion awareness module, only replacing warping based cost volume with sampling based one reduces end point error of $4.7\%$ on Sintel Final and $3.9\%$ on KITTI 2012. While only adding occlusion awareness module reduces error of $1.7\%$ on Sintel and $5.4\%$ on KITTI. It concludes that sampling is more remarkable on Sintel which contains more non-rigid motion and has higher requirement on distinct scene structure, while occlusion awareness is more effective on KITTI that involves large movement and occluded regions. Finally, combining these two approaches gets best results that reduces end point error of $6.4\%$ on Sintel Final and $11.0\%$ on KITTI 2012, demonstrating our contributions are cooperative and complementary.

To compare with state-of-the-art methods [5, 8, 6, 9, 12], we first evaluate above two stage trained OAS-Net on Sintel [16] and KITTI [17] training sets. For comparison on Sintel test benchmark, we fine-tune OAS-Net on mixed Sintel datasets, where we adopt batch size 4, with 2 from Clean, and 2 from Final. Similar to IRR-PWC [9], learning rate is initially set to $5e-5$ and disturbed every $150k$ iterations over total $600k$ iterations. To evaluate on KITTI 2012 test benchmark, we fine-tune OAS-Net on mixed KITTI 2012 and KITTI 2015 training sets following the same schedule on Sintel while reducing amplitude of spatial augmentation. All results together with model size and running speed (measured on Sintel resolution) are listed in Table 2.

OAS-Net performs better than all the other approaches on Sintel Clean test dataset, it improves about $5\%$ than the second one. OAS-Net also surpasses PWC-Net [6] on Sinel Final test set slightly, but is exceeded by IRR-PWC [9] and HD3-Flow [12]. However, our model has less learnable parameters, i.e., $6.16$M vs $6.36$M$/$$39.56$M and runs several times faster, i.e., $0.03$s vs $0.15$s$/$$0.08$s. We further evaluate it on a more challenging real-world KITTI dataset. As shown in Table 2, OAS-Net achieves the same best result as HD3-Flow [12] on KITTI 2012 test benchmark with end point error of $1.4$, which reduces $12.5\%$ error than LiteFlowNet [8].

To verify our approaches visually, we show one example on each Sintel Final test and KITTI training sets, as depicted in Fig 3 and Fig 4. It can be seen that our occlusion awareness maps have correctly emphasized the probable occluded regions, and help to improve optical flow accuracy, such as the left bottom corner of Fig 3 and fast moving edges in Fig 4.