PyNET-CA: Enhanced PyNET with Channel Attention for End-to-end Mobile Image Signal Processing

First, we define some of the modules that constitute the PyNET-CA model (Figure 2). The channel attention module (CA) of PyNET-CA follows from [19]. Global average pooling is first applied to the height and width dimension of the features to output a vector with length corresponding to the number of input channels. The pooled vector is linearly mapped, passed through the nonlinear ReLU activation, and is again linearly mapped to match the number of input channels. The length of output features after the first linear mapping is determined by the reduction ratio $r$, which reduces the number of channels by $\frac{1}{r}$ of the input channels. The final features are squashed to range $[0,1]$ by the sigmoid function, and are multiplied element-wise with the channels of the input to account for the level of attention for each channels.

The DoubleConv module is defined as two sequential operations of 2D convolution followed by a LeakyReLU activation. The kernels of the two convolution layers have the same shape, which is one of $3\times 3$, $5\times 5$, $7\times 7$, $9\times 9$. The input image or feature to the DoubleConv module is reflect-padded before the convolution to match the size of the output feature. Unlike the original PyNET [7], we apply instance normalisation after only the second convolution layer where needed.

The MultiConv channel attention (MCCA) module is the basic building block of our model. It comprise concatenating the features from the DoubleConv modules and a channel attention module. We specify four types of MCCA (MCCA3, MCCA5, MCCA7, MCCA9) based on the kernel size of the DoubleConv module. The MCCA3 module has only one $3\times 3$ DoubleConv module, while the MCCA5 module concatenates the output of $3\times 3$ and $5\times 5$ DoubleConv module, and so forth. Lastly, channel attention is applied to the concatenated features from the DoubleConv modules. The reduction ratio $r$ of the channel attention module is set to 1, 2, 3, 4 for MCCA3, MCCA5, MCCA7, MCCA9, respectively.

To account for both the global and local features of the image, PyNET-CA has an inverted pyramidal structure as in Figure 3. Given a RAW image with Bayer pattern $\mathbf{I}_{raw}\in\mathbb{R}^{H\times W}$, Bayer sampling function $f_{Bayer}$ is first applied to obtain

At each levels, the images are downscaled by $2\times 2$ max pooling followed by the MCCA3 module to serve as the input feature $F_{in}^{k}$ at level $k$,

We denote the set of operations at each level $k$ that process the input feature $F_{in}^{k}$ as function $H^{k}$. The function $H^{k}$ is composed of of a number of MCCA modules along with residual connections, within-level and between-level skip connections to compute the output feature $F_{out}^{k}$,

The composition of the operators of $H^{k}$ differs at each levels (Figure 3).

Lastly, the reconstructed RGB image $\mathbf{\hat{I}}_{rgb}^{k}$ is obtained by a $3\times 3$ convolution layer and tanh activation of the output features,

One important point regarding the inverted pyramidal structure lies in (5). It should be noted that the output feature at level $k$ is computed not only with the input feature at level $k$, but also with the output feature from lower-resolution level $k+1$ as a prior. Because the network is trained progressively, the prior $F_{out}^{k+1}$ contains meaningful information of the downsampled target image. The bilinear downscaling of the target image at each level corresponds to low-pass filtering, and explicitly emphasizes global feature information as the levels elevate. This information can be passed onto the lower levels effectively and recursively by concatenation.

Because of the Bayer sampling function in (1), the enhanced features need to be upsampled at the last level of the PyNET-CA. This is an ill-posed problem as in the case of the SISR problems. At the final level of the original PyNET structure, enhanced features are upsampled with bilinear interpolation or transposed convolution, and then convolved with a $3\times 3$ kernel convolution layer to output the final image. However, bilinear interpolation or transposed convolution can cause blurring or checkerboard artifacts in the upsampled image. Furthermore, reconstructing final image with the convolution layer after upsampling the image is computationally inefficient.

For computational efficiency and better image quality, the proposed PyNET-CA upsamples the image with the MCCA9 module, followed by $1\times 1$ convolution layer and upsamples the features by subpixel shuffling [14] at the final level of the model,

We denote the RGB reconstruction module (7) as the subpixel reconstruction module (SRM).

For training the network, we used the Zurich RAW to RGB (ZRR) dataset. One important fact about the dataset is that it consists of a RANSAC [4] aligned pair of RAW and RGB images, taken separately from a smartphone (Huawei P20) and a DSLR camera (Canon 5D Mark IV), respectively. Considering the collection process, the dataset is prone to significant pixel-level mismatch between the RAW and RGB image pairs if the object is moving or if the object has large area of reflection (see Figure 4). We screened out the images with large area of reflection (e.g. on cars, on windows) or moving objects (e.g. people, animal, vehicles on the road). There were 1,679 out of 46,839 (3.58%) training image pairs excluded from this screening, leaving out 45,160 image pairs for training the model. For testing, we used 1,204 image pairs provided by the AIM 2020 challenge organizers without excluding any image pairs. All RAW and RGB images were $448\times 448$ cropped patches, and were casted into single precision floating point data centered and scaled to the range $[-1,1]$ to stabilise the training.

The PyNET-CA model is progressively trained from the level with the lowest resolution. The target image $\mathbf{I}_{rgb}$ is downsampled with bilinear interpolation to match the resolution of the reconstructed images $\mathbf{\hat{I}}_{rgb}^{k}$ at each level. The loss function used for training the PyNET-CA is a linear combination of the mean squared error (MSE) loss, the perceptual loss using one VGG layer relu5_4, and the multi-scale structural similarity index measure (MS-SSIM) loss,

Basically the MSE loss is minimised at all levels with $\lambda_{1}$ set to 1.0, and other coefficients were adjusted with respect to the $\lambda_{1}$. For training level 5 and level 4, only the MSE loss is minimised, with $\lambda_{2}$ and $\lambda_{3}$ set to 0.0. From level 3, the perceptual loss is minimised to account for the perceptual similarity with coefficient $\lambda_{2}=0.01$, $\lambda_{3}=0.0$. At the last 0th level, the MS-SSIM loss is maximised (i.e., negatively minimised). Different MS-SSIM scaling coefficient is used for training the model based on the goal of the task. For the fidelity task, which is to achieve highest metric score on the peak signal-to-noise ratio (PSNR) and the MS-SSIM, the MS-SSIM scaling coefficient $\lambda_{3}$ is set to 0.01. For the perceptual task, on the other hand, we set the MS-SSIM scaling coefficient to 0.1.

The PyNET requires a long training time, especially at the high-resolution levels. To address this issue and hasten convergence, we employ the one-cycle policy of the learning rate for training the PyNET-CA [15]. At each levels, the learning rate starts with $5.0\times 10^{-5}$, reaches the maximum learning rate $1.0\times 10^{-4}$ at the early 20% of the training, and gradually decays to $5.0\times 10^{-7}$ by the end of the training as in Figure 5. We train the model for 16 epochs per level, which is around 68% decrement of the training epochs at the last level compared to [7].