Unrolling Plug-and-Play Network for Hyperspectral Unmixing

The plug-and-play framework benefits from variable splitting technique, allowing the utilization of denoising priors to effectively tackle a range of image restoration problems. It has been used for hyperspectral unmixing [29]. Typically, the denoising regularization is plugged for the abundances with (4) rewritten as:

We use ADMM to resolve the optimization problem described in (5). By introducing an auxiliary variable $\mathbf{V}$ to replace $\mathbf{A}$ in the regularization term of (5) and adding constraint $\mathbf{V}=\mathbf{A}$, the original problem is decomposed into easier and more manageable subproblems. The iterative optimization process of (5) is as follows:

in which $\mathbf{G}$ is the dual variable, and $\rho$ is the penalty parameter. The process involves two essential operators. The first is a blind unmixing operator to estimate endmembers and abundances. The second step can be viewed as a denoising of $\mathbf{A}+\mathbf{G}$. The regularization term $\mathcal{R}(\mathbf{V})$ can be implicitly coped by incorporating a denoising operator ($\mathsf{C}(\cdot)$), i.e.,

In general, the denoiser $\mathsf{C}(\cdot)$ can be any readily available denoising operator. For example, the well-known nonlocal means denoising (NLM) [30] and block-matching and 3D filtering (BM3D) [31] have been plugged to capture image priors and get high-quality unmixing results. This provides the possibility of integrating CNN-based denoisers with robust prior learning from mass image data and addressing the drawback of insufficient volume of hyperspectral data.

The deep unrolling paradigm involves unfolding iterative optimization algorithms into trainable deep architectures. One popular deep unrolling unmixing method is unfolding the CSR for sparse unmixing [27]. With known $\mathbf{M}$ and $\mathcal{R}(\mathbf{A})=\|\mathbf{A}\|_{1}$, it introduces an auxiliary variable $\mathbf{V}=\mathbf{A}$. The ADMM iteratively addresses the optimization problem through the following steps:

in which $\mathbf{W}=\left(\mathbf{M}^{\top}\mathbf{M}+\rho\mathbf{I}\right)^{-1}\mathbf{M}^{\top}$ and $\mathbf{Q}=\left(\mathbf{M}^{\top}\mathbf{M}+\rho\mathbf{I}\right)^{-1}\rho$. $\text{Soft}(a,\theta)=\text{sign}(a)(|a|-\theta)_{+}$ is the soft-threshold operator to resolve the $\ell_{1}$-norm. The iterative steps from (10) to (12) can be unfolded into network layers, denoted as $\text{Layer}_{A}$, $\text{Layer}_{V}$ and $\text{Layer}_{G}$ with learnable parameters $\left\{\mathbf{W}^{(k+1)},\mathbf{Q}^{(k+1)},\theta^{(k+1)}\right\}$, where $\theta^{(k+1)}$ replaces the role of $\frac{\lambda}{\rho}$. Thus (10) to (12) can be reexpressed as follows:

Each iteration is decomposed into a single network layer, and combining these layers many times is analogous to performing multiple iterations of CSR. The unrolling deep network can be trained by an end-to-end manner. By this way, the knowledge of classical iterative algorithms are infused in the design of deep network for hyperspectral unmixing. However, the networks obtained by unrolling iterative algorithms are restricted to one pixel only and not take into account of spatial structure information. By some modifications of networks, we can employ the well-designed layers such as convolutional networks to add more spatial priors.

In (4), handcrafting an effective regularizer $\mathcal{R}(\mathbf{M},\mathbf{A})$ and devising an efficient algorithm to solve the objective function are challenging tasks. Rather than using this intricate path, we propose to derive priors from various image data and integrate them into the optimization process by a plug-and-play strategy, where the regularization is implicitly designed by the utilization of a denoising algorithm. In hyperspectral unmixing task, the spatial information is embedded in the abundance, and the spectral information are contained in the endmember. Thus we plug regularization on abundance to fully exploit the spatial information. We introduce regularization by denoising (RED) as the regularizer to add denoising priors, and the corresponding hyperspectral unmixing optimization problem can be reformulated as:

where $\Omega_{\mathcal{D}}$ is an indicator function of a set $\mathcal{D}$ and is defined as:

$\mathcal{R}_{\text{RED}}(\mathbf{A})$ is defined as:

where $\mathsf{C}(\cdot)$ is an off-the-shelf denoiser. RED depends on the inner product between the solution $\mathbf{A}$ and its residual after denoising $\mathbf{A}-\mathsf{C}(\mathbf{A})$. Under mild assumptions, it demonstrates advantageous derivative properties with $\nabla\mathcal{R}(\mathbf{A})=\mathbf{A}-\mathsf{C}(\mathbf{A})$. This formulation also provides flexibility in accommodating any denoising engines like original plug-and-play denoising framework.

We use ADMM algorithm to solve (16). Two auxiliary variables $\mathbf{V}_{1}$ and $\mathbf{V}_{2}$ are introduced, and the objective function is rewritten as:

The associated augmented Lagrangian is expressed as:

where $\mathbf{G}_{1}$ and $\mathbf{G}_{2}$ are dual variables, and $\alpha$ and $\beta$ are nonnegative parameters. Then the ADMM algorithm involves solving four subproblems at each iteration. In our work, we unroll the ADMM algorithm into a deep network by associating each iteration with a single network layer and then stacking a finite number of layers together. Through unrolling, we can use back-propagation to learn regularization parameters that are difficult to handcraft. This approach allows us to capture image priors through a well-designed network structure and add physical interpretability to the deep network. The proposed unrolling layers (PnP-Net) are presented in detail as follows.

Each unrolling block represents a single iteration of the ADMM based plug-and-play unmixing method and consists of 6 parts: $\text{Layer}_{A}$, $\text{Layer}_{V_{1}}$, $\text{Layer}_{G_{1}}$, $\text{Layer}_{M}$, $\text{Layer}_{V_{2}}$ and $\text{Layer}_{G_{2}}$. As shown in Fig. 4, we use $K$ iteration blocks to conduct the architecture of our proposed method, which can mimic the ADMM based unmixing algorithm with $K$ iterations. For the $k$th block, the learnable parameters are $\bm{\Theta}^{(k)}=\left\{\mathbf{R}^{(k)},\mathbf{W}_{1}^{(k)},\mathbf{H}^{(k)},\mathbf{Q}_{1}^{(k)},\theta_{1}^{(k)},\theta_{2}^{(k)},\mathbf{W}_{2}^{(k)},\mathbf{Q}_{2}^{(k)},\theta_{3}^{(k)},\theta_{4}^{(k)}\right\}$. Each block in the network has its own set of learnable parameters that are not shared across blocks, which has demonstrated flexibility and strong learning capability for unmixing task.

Additionally, the PnP-Net architecture can be divided into two parts: the endmember estimation network and the abundance estimation network. Each block of the endmember estimation network contains the $\text{Layer}_{M}$, $\text{Layer}_{V_{2}}$ and $\text{Layer}_{G_{2}}$ components, there are $K$ blocks to estimate the endmember. The abundance estimation network comprises three parts, namely $\text{Layer}_{A}$, $\text{Layer}_{V_{1}}$ and $\text{Layer}_{G_{1}}$. $K$ layers are applied to estimate the abundance.

Our PnP-Net is trained by a blind manner with the input hyperspectral image $\mathbf{X}$. We use the mean square error (MSE) to compute the differences between the input and its reconstruction. Compared to calculating the MSE loss directly using the reconstruction of the final $K$th block output, calculating it for each block can better constrain the training process of the network. Thus, we use the weighted sum of the reconstructions from each block to compute the loss function, which is defined as:

where $\beta_{k}$ denotes the importance of each term, $\hat{\mathbf{X}}^{(k)}=\hat{\mathbf{M}}^{(k)}\hat{\mathbf{A}}^{(k)}$ is the corresponding reconstructed image of the $k$th block.

Parameter initialization plays an important role in our PnP-Net, which can accelerate the convergence speed and enhance the endmember and abundance estimation accuracy. The endmember matrix $\mathbf{M}^{(0)}$ and abundance matrix $\mathbf{A}^{(0)}$ are calculated by VCA [34] and FCLS [35]. Specially, we obtain the initialization of $\mathbf{W}_{1}$ and $\mathbf{Q}_{1}$ according to $\mathbf{A}^{(0)}$ and $\mathbf{M}^{(0)}$, and other values are initialized with zeros.

Our PnP-Net method is flexible to plug various deep denoisers. In this work, we choose 3 novel denoisers, i.e. DnCNN [36], IRCNN [37] and SCUNet [38], as examples. We pretrain these models using images of Waterloo Exploration Database [39], DIV2K [40], and Flick2K [41] datasets with different noise settings. It is worth noting that the parameters of denoiser are fixed during training which can reduce computational stress and take advantage of outer priors. Our PnP-Net with these denoisers are named as PnP-Net-1, PnP-Net-2 and PnP-Net-3. During training, we set the learning rate to $5\times 10^{-4}$. The number of iteration blocks ($K$) is set to 5. $L$ and $P$ are set to 3 with the kernel size setting to $1\times 1$, $3\times 3$ and $5\times 5$. We set $\{\beta_{k}\}_{k=1}^{K}$ to $\left\{1\times 10^{-4},1\times 10^{-3},1\times 10^{-2},1\times 10^{-1},1\right\}$.

We use several indicators to evaluate the unmixing results quantitatively in terms of estimated abundances, extracted endmembers and reconstructed images. In addition, we have no available ground truth for the real datasets, we only evaluate the reconstructed image quantitatively and subjectively evaluate the visual effect.

The root mean square error (RMSE) is applied to assess the differences between the ground truth and estimated abundances and endmembers. We respectively name them as aRMSE and mRMSE, which are defined as follows:

and

where $\hat{\mathbf{a}}_{i}$ and $\mathbf{a}_{i}$ represent the estimated and real abundances of the $i$th pixel, and $\hat{\mathbf{m}}_{r}$ and $\mathbf{m}_{r}$ denote the $r$th extracted and corresponding real endmembers.

We apply the mean spectral angle distance (SAD) to evaluate the degree of distortion between the real and reconstructed spectral signatures, and mSAD denotes the SAD between the estimated and real endmembers. They are calculated by:

and

where $\mathbf{y}_{i}$ and $\hat{\mathbf{y}}_{i}$ represent the real and reconstructed spectra.

We also use the peak signal-to-noise ratio (PSNR) to evaluate the differences between the real and reconstructed images, calculated by:

where $\mathrm{max}$ represents the highest pixel value in the reconstructed image.

We compare the proposed PnP-Net method with several conventional and state-of-the-art unmixing methods, i.e. SUnSAL-TV [5], gtvMBO [10], U-ADMM-BUNet [27], DIFCNN [28] and AERED [42]. To make a fair comparison, all methods are initialized by the endmembers extracted by VCA [34] and abundances estimated by FCLS [35]. All parameters of these methods are carefully chosen to achieve the best experimental results.

The first two compared methods are conventional methods. The SUnSAL-TV method uses the VCA to extract the endmembers, and a total variation spatial regularization is applied to constraint the estimated abundances. The gtvMBO method is a blind unmixing method with a data-driven graph total variation regularization, and its objective function is solved by ADMM. The U-ADMM-BUNet and DIFCNN are novel unrolling based unmixing methods. AERED integrates the autoencoder based unmixing network with RED, which combines the explicit and implicit priors. We name the AERED with NLM denoiser and block-matching and 4-D ffltering (BM4D) [43] deoiser as AERED-1 and AERED-2, respectively.

In order to quantitatively verify the unmixing results, we generate a set of synthetic data with the spatial size of $100\times 100$. We adapt the method described in [44] to produce these data. The ground-truth of abundance maps are shown in Fig. 5, and they are constrained to satisfy ASC and ANC. The abundance obtained by this strategy shows spatial correlation that similar to the real data. Four spectral signatures selected from the USGS spectral library with 224 bands are applied to generate the endmember matrix. The reference endmember spectra are shown in Fig. 6. Zero mean Gaussian with four levels of signal-to-noise-ratio (SNR), i.e. 5dB, 10dB, 20dB and 30dB, are added to the data.

Table I shows the aRMSE and PSNR results of the synthetic data, and Table II lists the mRMSE and mSAD results. The bold numbers in Tables I and II denote the best results, and the underlined numbers indicate the second best results. From the results of these two tables, it can be observed that the proposed PnP-Net unmixing framework achieves the best unmixing results compared to the benchmark methods. Compared to traditional unmixing methods such as SUnSAL-TV and gtvMBO, the proposed method eliminates the need for penalty parameter selection, instead utilizing a data-driven learning approach. In comparison to unrolling unmixing method, i.e. U-ADMM-BUNet and DIFCNN, this proposed PnP-Net method employs dynamic convolution to achieve multiscale information learning and fusion and gets better results. Our PnP-Net method leverages an externally pretrained denoiser on large datasets to provide outer priors, and also integrate a hyperspectral data-driven and physically interpretable neural network to learn internal priors. This combination results in superior unmixing performance. We can also observe that, due to the use of deep denoisers and deep neural networks, our proposed unmixing method is robust to noise.

Fig. 5 shows the abundance maps of the proposed method and the comparison methods for the synthetic data at 10dB. We observe that all these estimates are very close to the ground-truth. Moreover, it can be seen that the noise in the abundance maps of PnP-Net and AERED is obviously smaller than that of other comparison methods, which indicates the superiority of using denoiser to bring prior information. Furthermore, Fig. 6 illustrates the extracted endmembers of the proposed method. We observe that the result of our method is close to the ground-truth. Fig. 7 illustrates how the number of blocks affect the unmixing performance and also shows that how the number of kernels $L$ and $P$ impact the unmixing results.

We also evaluate our proposed method on real datasets. We conduct experiments on two widely used hyperspectral images obtained by airborne sensors, namely Jasper Ridge dataset and Muufl Gulfport dataset. The descriptions and experimental results about these two datasets are presented below.

The Jasper Ridge dataset is captured by Analytical Imaging and Geophysics (AIG) in 1999. It contains 224 bands covering the spectral range from 380nm to 2500nm. After removing the bands affected by water vapor and atmospheric (1-3, 108-112, 154-166, and 220-224), 198 bands are remained. The original size of this data is $512\times 614$, a popular region of interest with $100\times 100$ pixels are cropped. The four endmembers in this data are “water”, “soil”, “tree” and “road”.

The Muufl Gulfport dataset was captured by the CASI-1500 hyperspectral sensor over the University of Southern Mississippi Gulf Park Campus. The original image contains 72 bands with $325\times 220$ pixels. We remove the first four and the last four bands of the image, which contain a lot of noise. A subimage with the size of $130\times 90$ pixels is applied in our experiments. We reference the scene label ground-truth map from manually labeling in [45]. There are 5 pure materials in this subimage, namely “roof”, “grass”, “tree”, “shadow” and “asphalt”.

The visual results of the estimated abundance maps of these two datasets are shown in Fig. 8 and Fig. 9. The first four rows of Fig. 8 report the comparisons of our method and compared methods with respect to the Jasper Ridge dataset. All these methods decompose this data with four clear abundance maps. But we observe that the abundance maps of “road” material of SUnSAL-TV and gtvMBO is lighter than other methods. The reason may be that the data contains spectral variability, which makes these two conventional methods hard to extract representative endmember. Our PnP-Net gets clearer abundance results. The last line of Fig. 8 illustrates the reconstructed error between the reconstructed image and real image of these methods. We observe that our method gets good reconstruct results which are also consistent with the results in Table III. To some extent, it proves the advantages of our method to unmix the real data. The abundance maps of Muufl Gulfport dataset are illustrated in Fig. 9. It can be seen that the results of gtvMBO are noisy than other methods, especially for the “asphalt” endmember. The proposed PnP-Net method can estimate the abundance maps much clearer than the compared method. Our results also have more detailed spatial information and sharper edges. The last row of Fig. 9 presents the reconstructed error maps of all methods. DIFCNN shows higher results. It may cause by the loss function selected to train the model. It use the cross entropy-loss to train the network, while other compared deep methods apply the MSE-like loss to train the model. The latter is more conductive to reduce the Euclidean distance between the reconstructed and original image. Table IV lists quantitative results of the Muufl Gulfport data. We get that our method can obtain good reconstructions.