Transmission-Guided Bayesian Generative Model for Smoke Segmentation

Let us assume the noise-aware model prediction as:

where $f_{\theta}(x_{i},z_{i})$ is the prediction of a neural network, $z_{i}$ is a latent variable, $\epsilon_{i}$ is additive noise drawn from a normal distribution $\mathcal{N}(0,\sigma^{2}\mathbb{I})$, with $\sigma$ as standard deviation. Given input $x_{i}$ and latent variable $z_{i}$, network output will be corrupted by additive noise $\epsilon_{i}$, which we assume is conditioned on the input image. It will lead the network to force $\sigma$ to be small to reduce the influence of the noise in the prediction. Then, the latent variable $z_{i}$ can capture unobserved stochastic features that can affect the network’s prediction.

Following conventional practice of variational auto-encoders, we train the latent variable by maximizing the expected variational lower bound (ELBO) as:

where $P_{\omega}(z|X)$ is the posterior distribution of the latent variable $z$ or the inference model, and $P(z)$ is its prior distribution, which we defined as a standard Gaussian distribution $\mathcal{N}(0,1)$. $D_{KL}(P_{\omega}(z|x)||P(z))$ is the Kullback-Leibler divergence regularizer. We design the inference model $P_{\omega}(z|X)$ with five convolutional layers that maps input $x$ into a low dimensional latent variable $z$ to capture the inherent ambiguities of human labeling. Specifically, the output of the inference model is $\mu^{post}$ and $\sigma^{post}$, leading to $P_{\omega}(z|x)=\mathcal{N}(\mu^{post},\sigma^{post})$ where $\mu^{post}$ and $\sigma^{post}$ represent the mean and the standard deviation of the posterior distribution of $z$, and $z$ is then obtained with the reparameteration trick: $z=\mu^{post}+\sigma^{post}*\epsilon$, where $\epsilon\sim\mathcal{N}(0,1)$.

As $z$ can model the aleatoric uncertainty, and $P(\theta|D)$ can model epistemic uncertainty, with a Bayesian latent variable model, we aim to model both epistemic and aleatoric uncertainty, leading to predictive uncertainty estimation. We design our BNN based latent variable model (BNN-LVM) with a ResNet50 backbone (He et al. 2015) as encoder to obtain feature maps $\{f_{k}\}_{k=1}^{4}$ from different stages of the backbone. The latent variable $z$ is introduced to the network via two steps. Firstly, we sample $z$ from $P_{\omega}(z|x)$, and tile it to achieve the same spatial size as $f_{4}$. Secondly, we concatenate the tiled $z$ with $f_{4}$, and feed it to a $3\times 3$ convolutional layer to obtain $f_{4}^{\prime}$ of the same size as $f_{4}$, which will serve as the new backbone feature. To capture the epistemic uncertainty, the produced new features $\{\{f_{k}\}_{k=1}^{3},f_{4}^{\prime}\}$ are fed into four different $3\times 3$ convolutional layer with a dropout rate of 0.3 to get the new backbone features $\{s_{k}\}_{k=1}^{4}$. Subsequently, we use four residual channel attention modules (Fu et al. 2018) after each $s_{k}$ to obtain discriminative feature representations. Finally, we use a DenseASPP (Yang et al. 2018) module to aggregate higher-lower level features with enlarged receptive field and get the final prediction of our BNN-LVM as $s=P_{\theta}(y|x,z)$.

Given the model prediction $P_{\theta}(y|x,z)$ from our BNN module, the predictive (total) uncertainty can be measured as entropy of the mean prediction, which can be formulated as $\mathcal{H}[\mathbb{E}_{P(\theta|D),z\sim P_{\omega}(z|x)}[P_{\theta}(y|x,z)]]$ where $\mathcal{H}$ denotes the entropy and $\mathbb{E}$ denotes the expectation. The corresponding expectation can be approximated by Monte Carlo (MC) integration, leading to mean prediction as:

where $B$ is the iterations of MC sampling, and $\theta_{b}$ is the parameter set of the $b^{th}$ iteration of sampling. The predictive uncertainty is then defined as: $U_{p}=\mathcal{H}[p_{\mu}]$. The aleatoric uncertainty is the average entropy for a fixed set of model weights, which can be formulated as: $U_{a}=\mathbb{E}_{P(\theta|D),z\sim P_{\omega}(z|x)}[\mathcal{H}[P_{\theta}(y|x,z)]]$. The epistemic uncertainty is then defined as the gap between predictive uncertainty and aleatoric uncertainty: $U_{e}=U_{p}-U_{a}$.

With the loss function in Eq. 2 and the uncertainty estimation related loss in Eq. 6, we can already train our Bayesian latent variable model for smoke segmentation. However, we find that model with the smoke segmentation related loss in Eq. 4 fails to predict distinct boundaries at the ambiguous smoke edges. Inspired by image dehazing, we propose a transmission-guided local coherence loss, which exploits transmission as an essential feature to discriminate the smoke boundary. Our loss is based on the assumption that smoke objects tend to have a different transmission from most regions of the background. It can enforce similar predictions for pixels with similar transmission and close distance. Further, outdoor smoke images often suffer from quality degradation, i.e. low contrast. (Li et al. 2021b) shows that image quality degradation regions usually have lower transmission values. Inspired by (Li et al. 2021b), we use the reversed transmission value as the weight of our loss to encourage the model to focus on the degraded regions, which is defined as:

where $K_{m}$ is a $k\times k$ kernel centered at pixel $m$, $D(m,n)$ is the $L1$ distance between the prediction of pixels $m$ and $n$, and $W(m,n)$ is a modified bilateral kernel in the form:

where $\frac{1}{w}$ is the normalized weight term, $P(m),P(n)$, and $T(m),T(n)$ represent the spatial information and transmission information of the pixels $m$ and $n$ respectively, and $\sigma_{P}^{2},\sigma_{T}^{2}$ are the kernel bandwidth.

To estimate medium transmission, we calculate transmission information of each pixel of the intensity image by:

where $A$ is global atmospheric light, $I$ is the intensity image, $c$ is the color channel, $K(m)$ is a local patch of size $k\times k$ centered at pixel $m$ and then $\min\limits_{c}(\min\limits_{n\in K(m)}(\frac{I^{c}(n)}{A^{c}})$ is the normalized haze map defined at (He, Sun, and Tang 2011). We can see the estimated medium transmission is based on the global atmospheric light $A$. To get $A$, we follow (He, Sun, and Tang 2011) and obtain it by picking the top 0.1% brightest pixels in the dark channel of the intensity image.

So far, we have $\mathcal{L}(\theta,\omega;x)$ in Eq. 2 for the Bayesian latent variable model, $\mathcal{L}_{un}$ in Eq. 6 for the uncertainty estimation module, and $\mathcal{L}_{trans}$ in Eq. 7 to force the network focus on the image degradation regions. We further adapt entropy loss as a regularizer to encourage the prediction to be binary. Entropy loss is defined as the entropy of model prediction, which is minimized when the model produces binary predictions. However, we find that directly applying entropy loss to our task leads to overconfident predictions, causing extra false positives. To reduce the “binarization regularizer” for less confident regions, we add a temperature scaling term (Guo et al. 2017) to the original entropy loss, which uses learned total uncertainty as the temperature so that the model can produce softened predictions on high uncertainty pixels. Note that the temperature in our method is learned by the uncertainty estimation network rather than being a user-defined temperature as used in traditional temperature scaling methods. Our uncertainty-based temperature scaling can also be used to calibrate our model. A well-calibrated model cannot only reduce the gap between prediction confidence and model accuracy but also achieve better performance.

Finally, our total loss for BNN-LVM is defined as:

and empirically we set $\lambda_{1}=0.3$ and $\lambda_{2}=0.01$. With Eq. 11 we can update our Bayesian latent variable model, and the uncertainty estimation model $g_{\gamma}$ is updated with loss function $\mathcal{L}_{un}$ in Eq. 6.

We create our dataset based on an open wildfire smoke dataset (Zhou et al. 2016), the current largest synthetic smoke dataset (Yuan et al. 2018) and a few images from the internet. Due to the large redundancy of the synthetic dataset, we only choose 4,000 images from the synthetic dataset and discard images that contain duplicated smoke foreground. For real datasets, we choose 1,360 good quality smoke images from the (Zhou et al. 2016) dataset, and discard smoke images of poor quality, and we also choose 40 smoke images from the internet to further increase the diversity of the test set. As the real smoke datasets do not have pixel-wise annotations, we then annotate the dataset with both binary mask and scribble annotation. Finally, we end up with 5,400 images, where 400 real smoke images from it is used for testing, which is defined as “Total”, and the remaining 5,000 images are for training. To better explore the difficulty of different attributes in the smoke segmentation models, we further extract 100 images from the 400 test images, denoted as “Difficult”, which mainly contains transparent smoke and smoke with more diverse shapes.

We have carried out two experiments to verify the superiority of our dataset. Firstly, we train the same segmentation model on both SYN70K and our dataset SMOKE5K. It takes about three days to train on SYN70K and about 12 hours to train on our dataset. Then, we evaluate the two models on three synthetic datasets (DS01, DS02, DS03) from (Yuan et al. 2018), and the results are shown in Fig. 3 (a). We observe that although SYN70k is at least 14 times larger than our dataset, the model trained with our dataset still produces the lower mean square error (MSE) on all three synthetic test sets. Further, we evaluate the two models on our real image test set, and the result is shown in Fig. 3 (b), which further explains the effectiveness of our new training dataset.

Following previous smoke segmentation literature, we choose the Mean Square Error (Mse) ($\mathcal{M}$) as the evaluation metric. The average Mses (mMse) on the test set is used to measure performance. Mse computes the per-pixel accuracy of model prediction, which depends on the size of the smoke foreground, leading to inherent small Mse for smoke image with small size of smoke. To tackle the problem, we adopt the F-measure ($F_{\beta}$) as complementary to mMse for comprehensive performance evaluation, which is a weighted combination of precision and recall values. Further, as an uncertainty estimation model, we also adopt the reliability diagram (Guo et al. 2017) and dense calibration measure (ECE) (Zhang et al. 2020a) to represent model’s calibration degree.

In Fig. 5, we choose three representative images according to the difficulty to segment, denoted as simple, normal and difficult examples. We visualize our model predictions with the corresponding three types of uncertainties, namely “Aleatoric”, “Epistemic” and “Total” uncertainty. For the simple example (the first row), epistemic and aleatoric uncertainty capture similar patterns. For the medium example (the second row), epistemic uncertainty is increased on the translucent region of the smoke, which accounts for the out-of-distribution example of our training sets. For the difficult example (the third row), epistemic uncertainty is especially increased in fog regions. Further, we can see the total uncertainty captures both aleatoric uncertainty and epistemic uncertainty. In summary, we find that it is important to model aleatoric uncertainty for: (1) modeling ambiguity in annotations, and (2) modeling the non-rigid shape of objects. Further, it is important to model epistemic uncertainty for: (1) safety-critical applications, identifying out-of-distribution samples; and (2) the false-positive problems, where the epistemic uncertainty is usually increased in false-positive regions. The total uncertainty in general can provide more information about model prediction, leading to explainable and well-calibrated models.

We perform the following experiments to analyze each component of our model. All experiments are evaluated on the test set of SMOKE5K. To evaluate the effectiveness of our method, we set our baseline by removing the transmission loss and the uncertainty calibration entropy loss, which is shown as “1” in Table 4. Note that it means our baseline does not exploit estimated uncertainty. Then, we combine the transmission loss and uncertainty calibration entropy loss gradually, denoted in “2” and “4” in Table 4.

Further, as an uncertainty-aware learning method, we aim to achieve well calibrated model where model accuracy is consistent with model confidence. To measure the calibration quality, we follow the calibration definition in (Guo et al. 2017), and report the reliability diagrams in Fig. 6. We show the reliability diagrams of our model without calibration (setting 2 in Table 4) and our model after calibration (setting 4 in Table 4) in the second and third diagrams of Fig. 6, respectively. We also show reliability diagrams of UCNet (Zhang et al. 2021a) (another uncertainty estimation method using CVAEs (Sohn, Lee, and Yan 2015) for saliency detection) in the first diagram of Fig. 6. It can be seen that the gap between accuracy and confidence for UCNet is smaller than our model without calibration, leading to the ECE of 0.0076, but the calibration quality is still very low. With our uncertainty calibration entropy loss, the gap between confidence and accuracy for our model is significantly reduced, the dense calibration measure (ECE) is also reduced from 0.0159 to 0.0039, leading to a better calibrated model.

In this section, we introduce more detail about our SMOKE5K dataset. Different from the current largest synthetic dataset SYN70K, the real data in SMOKE5K is specifically for wildfire smoke captured from long-distance tower cameras, which is a cheap and effective option for monitoring wildfire. We mainly recognize six challenging object attributes of our dataset that are different from SYN70K and other conventional dense prediction tasks.

• •

  Similar background: The smoke object has a similar appearance with its background due to its translucent appearance, causing large ambiguity in both segmentation and labeling.

• •

  Similar foreground: The outdoor images often contain similar confusing objects (false-positive) like haze, cloud, fog, and sun reflection, which makes it difficult to identify smoke even by trained wildfire detection experts.

• •

  Object occlusion: The object can be partially occluded, resulting in disconnected parts.

• •

  Diverse shapes: Smoke has diverse shapes due to its non-rigid shape, increasing the difficulty to perform precise segmentation.

• •

  Small objects: The ratio between the smoke and the whole image can be even smaller than 1% as we want to detect the smoke in the early phase.

• •

  Diverse Locations: Most images in SYN70K (see Fig.7 (a)) dataset have strong center bias, where the smoke usually appears in the center of the image. Different from them, the smoke in our dataset can be in any position in the image.

Also, smoke images are difficult to precisely label due to the challenging attributes mentioned above, thus we also provide scribble annotations for weakly supervised smoke segmentation, shown in Fig.7 (e). Weakly supervised learning allows a model to learn from a weak supervision signal i.e. image-level (Zhou et al. 2015; Wang et al. 2018), scribble-level (Lin et al. 2016a), bounding box-level (Khoreva et al. 2016; Dai, He, and Sun 2015) annotations. Scribble annotation is especially suitable for smoke segmentation as: 1) compared with pixel-wise annotation, scribble annotation is much cheaper and faster (It only takes 2$\sim$5 seconds for each image); 2) compared with image-level annotation, scribble annotation is effective for localizing small objects; 3) compared with bounding box-level annotation, scribble annotation is more flexible for diverse shape objects.

The qualitative comparisons between SYN70K and our real images are visualized in Fig.7, which further demonstrates the superior of our dataset.