A Dual-Stage Deep Learning Framework for Pneumothorax Detection and Segmentation in Chest X-rays

Effective treatment of pneumothorax, a potentially deadly condition caused by air accumulation in the pleural cavity between the lung and chest wall, necessitates fast and accurate diagnosis. Chest radiographs continue to be the principal diagnostic technique for identifying pneumothorax, but being widely available, they can be difficult to interpret, even for skilled radiologists [1, 2]. In the study by Pedrosa et al. [3], the importance of prompt diagnosis and the mounting strain on healthcare systems around the world have prompted research into automated diagnostic solutions, especially those that make use of deep learning and artificial intelligence.

In the last few years, automatic pneumothorax detection has progressed dramatically, ranging from basic computer vision techniques to complex deep learning architectures. Urooj et al. [4] presented the potential of Deep Convolutional Neural Networks (DCNN) enhanced with Network in Network (NIN) preprocessing, achieving AUC of 0.9844. This foundation was more upon by incorporating transfer learning approaches, as evidenced by developments in segmentation techniques. Bondarenko and Syryh [5] represented an improved UNet++ architecture that achieved a Dice similarity coefficient of 88.31% and IOU of 83.1%, marking a significant improvement in the precision of pneumothorax localization.

In order to get better local and global information in chest radiographs, Sanati et al. [1] supposed a combination for convolutional and transformer network. By integrating U-Net segmentation with DenseNet structures, Manikandan et al. [2] proposed this hybrid technique with increasing the accuracy to 98.5 percent. frameworks called Lesion-aware categorization were recently introduced by Pedrosa et al. [3] highlighting the significance of interpretability in clinical situations.

While Urooj et al. [4] demonstrated the effectiveness of computer-assisted systems by acquiring 97.77 percent accuracy by DenseNet-169. Bondarenko and Syryh [5] enhance the field by providing approaches for predicting illness development. Kancherla et al. [6] improved segmentation accuracy with residual connections; Deng et al. [7] provided a way to accurately identify pneumothorax illnesses by testing different models and hyperparameter configurations.

Although these significant improvements, the current pneumothorax detection method still has substantial shortcomings. While individual models have demonstrated impressive performance in either segmentation or classification tasks, frameworks that effectively integrate both abilities are clearly lacking. Current methods often fail to utilize the supplementary data offered by both raw images and segmentation masks in their final diagnostic decisions.

Our suggested dual-stage strategy to solve the shortcomings of current methods, is motivated by these gaps in the present research environment.

The proposed approach offers a two-phase strategy for classifying and detecting pneumothorax in chest X-ray images. There are two main parts to the framework:

(1) U-Net based segmentation network that uses various backbone architectures to figure out the pneumothorax's spatial extent.

(2) Hybrid Classification System combines deep learning feature extraction with gradient enhancement classification.

Clinicians can obtain both localization and diagnostic information with this integrated technique and that enables both pixel-wise detection of pneumothorax zones and overall diagnostic categorization. Through preprocessing processes, unique loss functions, and comprehensive data augmentation strategies to handle the fundamental complications of chest X-ray interpretation, the methodology solves challenges of medical image analysis.

3.1 Problem formulation

It is essential to formulate the complicated computer vision problem of automatically detecting and classifying pneumothorax in chest radiographs as an integrated optimization problem, which includes binary classification and semantic segmentation. This dual-objective challenge necessitates both pixel-by-pixel detection of pneumothorax and a global assessment of its occurrence.

Let $\mathcal{D}=\left\{\left(\mathrm{I}_{\mathrm{i}}, \mathrm{M}_{\mathrm{i}}, \mathrm{y}_{\mathrm{i}}\right)\right\}_{\mathrm{i}=1}^{\mathrm{N}}$ represent dataset of N image samples, where:

- $I_i \in \mathbb{R}^{H \times W}$ denotes the input chest X-ray image

- $M_i \in\{0,1\}^{H \times W}$ represents the binary segmentation mask

- $y_i \in\{0,1\}$ indicates the presence of pneumothorax Our objective is to learn two mapping functions:

- Segmentation Function: $f_\theta: \mathbb{R}^{H \times W} \rightarrow[0,1]^{H \times W}$ where $\theta$ represents the learnable parameters of the segmentation network.

- Classification Function: $g_\phi: \mathbb{R}^{H \times W} \times[0,1]^{H \times W} \rightarrow [0,1]$ where $\phi$ represents the learnable parameters of the classification network.

The optimization problem can be formulated as:

$\min _{\theta, \phi} \mathcal{L}_{s e g}\left(f_\theta(I), M\right)+\lambda \mathcal{L}_{c l s}\left(g_\phi\left(I, f_0(I)\right), y\right)$         (1)

where, λ is a weighting parameter balancing the two objectives.

This formulation encapsulates the interdependence between segmentation and classification tasks, where the segmentation output $f_\theta(I)$ serves as an additional input to the classification function $g_\phi$, allowing the classifier to leverage both global image features and localized pneumothorax information. The joint optimization ensures that the segmentation network learns to identify relevant regions that contribute to accurate classification while maintaining pixel-wise accuracy in delineating pneumothorax boundaries.

3.2 Theoretical framework

The proposed solution is grounded in a comprehensive theoretical framework that integrates concepts from deep learning, computer vision, and statistical learning theory. Our approach builds upon the Universal Approximation Theorem for neural networks, which guarantees that sufficiently wide neural networks can approximate any continuous function on compact subsets of $\mathbb{R}^n$. This theoretical foundation is particularly relevant for medical image analysis, where complex patterns must be learned from high-dimensional data while maintaining robustness and generalizability.

Let $\mathcal{X}=\mathbb{R}^{\mathrm{H} \times \mathrm{W}}$ be our input space and $\mathcal{Y}= \{0,1\}^{\mathrm{H} \times \mathrm{W}} \times\{0,1\}$ be our output space. The learning problem can be viewed as finding functions in the hypothesis space $\mathcal{H}$ that minimize the expected risk:

$\mathcal{R}(h)=\mathbb{E}_{(x, y) \sim \mathcal{D}}[\mathcal{L}(h(x), y)]$          (2)

where, $\mathrm{h} \in \mathcal{H}$ and $\mathcal{L}$ is our composite loss function. Given the finite nature of our training set, we minimize the empirical risk while maintaining generalization through appropriate regularization techniques.

3.2.1 Data image preprocessing

The preprocessing stage is crucial for ensuring optimal model performance and convergence. The experiments were conducted using the SIIM-ACR Pneumothorax Segmentation Dataset, containing 12,047 annotated chest X-ray images [26]. We establish a formal preprocessing framework that addresses the specific challenges of medical image analysis, including intensity normalization, spatial standardization, and data augmentation for improved generalization.

Define the preprocessing pipeline $\mathcal{P}$ as a composition of transformations:

The preprocessing methodology is designed to satisfy several key theoretical properties:

Invariance Properties:

- Translation invariance: $\mathcal{P}\left(T_{\Delta}(I)\right) \approx \mathcal{P}(I)$ for small translations $\Delta$

- Scale invariance: $\mathcal{P}\left(S_\alpha(I)\right) \approx \mathcal{P}(I)$ for scale factors $\alpha$ near 1

- Rotation invariance: $\mathcal{P}\left(R_\theta(I)\right) \approx \mathcal{P}(I)$ for small angles $\theta$

Statistical Properties:

- Normalized intensity distribution: $\mathbb{E}[\mathcal{P}(I)]=0$, $\operatorname{Var}[\mathcal{P}(I)]=1$

- Preserved spatial correlations for anatomical structures

- Controlled noise characteristics

Information Preservation:

- Minimal loss of diagnostically relevant information

- Preservation of edge and texture information critical for pneumothorax detection

- Maintenance of spatial relationships between anatomical structures

$\mathcal{P}=T_n \circ T_{\text {aug }} \circ T_r$             (3)

where:

• $T_r$ is the resizing transformation.

• $T_{\text {aug }}$ represents the augmentation function defined as.

• $T_n$ is the normalization function.

3.3 Segmentation architecture

The segmentation component of our framework is built upon deep convolutional neural networks, specifically leveraging the U-Net architecture with various backbone networks. The architectural design is motivated by the need to capture both fine-grained local features crucial for boundary detection and global contextual information necessary for understanding anatomical relationships in chest X-rays. Our framework extends the traditional U-Net architecture by incorporating modern deep learning advances and domain-specific optimizations for medical image analysis.

Let $\mathcal{F}=\left\{\mathrm{f}_\theta: \mathcal{X} \rightarrow \mathcal{Y} \mid \theta \in \Theta\right\}$  represent the class of learnable segmentation functions, where $\Theta$  denotes the parameter space. The segmentation network implements a mapping:

$f_\theta: \mathbb{R}^{H \times W} \rightarrow[0,1]^{H \times W}$        (4)

Through a sequence of nested function compositions that form the encoder-decoder architecture with skip connections:

$f_\theta=g_d \circ h_{\text {skip }} \circ g_e$        (5)

where, $g_e$ is the encoder pathway, $g_d$ is the decoder pathway, and $h_{\text {skip}}$ represents the skip connection mechanism.

3.3.1 U-Net framework

The U-Net architecture implements a multi-scale feature extraction and synthesis framework that can be formalized as follows:

• Encoder Pathway: Let $E_k$ represent the encoding operation at level k:

$E_k: \mathbb{R}^{C_k \times H_k \times W_k} \rightarrow \mathbb{R}^{C_{k+1} \times H_{k+1} \times W_{k+1}}$        (6)

where the feature transformation at each level is defined as:

$\begin{aligned} & E_k(X) =\operatorname{Pool}\left(B N\left(\sigma\left(\operatorname{Conv}_{k, 2}\left(B N\left(\sigma\left(\operatorname{Conv}_{k, 1}(X)\right)\right)\right)\right)\right)\right)\end{aligned}$           (7)

with:

• $\operatorname{Conv}_{k, i}$ representing the $i$-th convolution at level k

• BN denotes batch normalization

• σ being the ReLU activation function

• Pool implements max pooling with stride 2

• Skip Connections: The skip connection mechanism $S_k$ at level k is defined as:

$S_k:\left(X_{d, k}, X_{e, k}\right) \mapsto \operatorname{Concat}\left(X_{d, k}, \operatorname{Crop}\left(X_{e, k}\right)\right)$           (8)

where, $X_{d, k}$ and $X_{e, k}$  are decoder and encoder features respectively.

• Decoder Pathway: The decoder operation $D_k$ at level k can be expressed as:

$\begin{gathered}D_k: \mathbb{R}^{C_{k+1} \times H_{k+1} \times W_{k+1}} \times \mathbb{R}^{C_k \times H_k \times W_k} \rightarrow  \mathbb{R}^{C_k \times H_k \times W_k}\end{gathered}$           (9)

• Multi-Scale Feature Integration: The final prediction integrates features across all scales:

$\widehat{Y}=\sigma\left(\operatorname{Conv}_{1 \times 1}\left(\sum_{k=1}^K w_k F_k\right)\right)$            (10)

where:

• $F_k$ represents features at scale k

• w_k are learnable scale weights

• $\operatorname{Conv}_{1 \times 1}$  is the final projection to output space

The U-Net architecture implements an encoder-decoder framework with skip connections. For an input tensor $\mathrm{X} \in \mathbb{R}^{\mathrm{C} \times \mathrm{H} \times \mathrm{W}}$, the encoder pathway E consists of k blocks:

$E_k(X)=\operatorname{Pool}\left(\sigma\left(\operatorname{Conv}_k\left(\sigma\left(\operatorname{Conv}_k(X)\right)\right)\right)\right)$         (11)

where, $Conv_k$ represents convolutional layers at level k.

3.3.2 Loss function formulation

The optimization of our segmentation network is guided by a carefully designed composite loss function that addresses multiple aspects of the segmentation task. Given the inherent challenges in medical image segmentation, including class imbalance, varying region sizes, and the critical importance of boundary accuracy, we propose a weighted combination of complementary loss terms. Each component of our loss function is designed to address specific challenges in pneumothorax segmentation while maintaining stable training dynamics.

The combined loss function $\mathcal{L}_{\text {total }}$ is defined as:

$L_{\text {total }}=\alpha L_{B C E}+\beta \mathcal{L}_{\text {Dice }}+\gamma \mathcal{L}_{\text {Focal }}$          (12)

where:

$\mathcal{L}_{B C E}=-\frac{1}{N} \sum_{i=1}^N\left[y_i \log \left(\hat{y}_i\right)+\left(1-y_i\right) \log \left(1-\hat{y}_i\right)\right]$             (13)

This term provides pixel-wise supervision and encourages probabilistic prediction calibration.

$\mathcal{L}_{\text {Dice }}=1-\frac{2|X \cap Y|+\epsilon}{|X|+|Y|+\epsilon}$          (14)

where, $\epsilon=1 \times 10^{-5}$ is a smoothing factor. The Dice loss contributes:

• Region-based optimization: Focus on spatial overlap

• Scale invariance: Performance independent of region size

• Gradient characteristics: $\frac{\partial \mathcal{L}_{\text {Dice }}}{\partial X} \propto \frac{Y}{(|X|+|Y|)^2}$

• Focal Loss:

$\mathcal{L}_{\text {Focal }}=-\alpha_t\left(1-p_t\right)^\gamma \log \left(p_t\right)$          (15)

where, $p_t$ is the predicted probability and $\gamma=2$ is the focusing parameter. This component provides:

• Dynamic weighting: $\mathrm{w}\left(\mathrm{p}_{\mathrm{t}}\right)=\left(1-\mathrm{p}_{\mathrm{t}}\right)^\gamma$

• Class balancing: Through $\alpha_t$ parameter

• Gradient modulation: $\frac{\partial \mathcal{L}_{\text {Focal }}}{\partial \mathrm{p}_{\mathrm{t}}} \propto\left(1-\mathrm{p}_{\mathrm{t}}\right)^{\gamma-1}$

The weights $\alpha, \beta$, and $\gamma$ are determined through empirical validation to optimize the trade-off between different loss components.

3.4 Classification framework

3.4.1 Feature extraction

The feature extraction stage implements a dual-stream architecture that processes both the original chest X-ray images and their corresponding segmentation masks. This approach enables the model to leverage both global image characteristics and localized pneumothorax indicators.

The feature extraction process can be formalized as:

$\Phi(I, M)=\left[\phi_{-} C N N(I) ; \phi_{-} C N N(M)\right]$         (16)

where, $\phi_{C N N}$ represents the DenseNet feature extractor and [;] denotes concatenation.

The DenseNet feature extractor implements dense connectivity patterns defined by:

$x_l=H_l\left(\left[x_0, x_1, \ldots, x_{l-1}\right]\right)$        (17)

where, $H_l$ represents the composite function of operations:

3.4.2 XGBoost classification

The XGBoost classifier extends traditional gradient boosting by incorporating second-order gradients and regularization. The model employs additive training to combine weak learners.

$\hat{y}_i=\sum_{k=1}^K f_k\left(x_i\right), f_k \in \mathcal{F}$          (18)

where, $\mathcal{F}$ is the space of regression trees.

3.5 Training dynamics

Adaptive optimization strategy implements by the training process a with dynamic learning rate adjustment and early holding techniques to avoid overfitting and guarantee strong convergence.

3.5.1 Learning rate schedule

Implements a plateau-based reduction strategy with momentum correction. The learning rate $\eta_t$ at epoch t follows:

$\eta_t=\eta_0 \cdot \prod_{i=1}^k \alpha_i$          (19)

$\alpha_i=0.5$ if plateau is detected, k is the number of learning rate reductions.

3.5.2 Early stopping criterion

Mechanism implements a patience-based approach with validation loss monitoring. Training terminates when:

$\mathcal{L}_{\text {val }}^t>\min _{i \in[t-p, t-1]} \mathcal{L}_{\text {val }}^i$          (20)

where, $\mathcal{L}_{\text {val }}^t$ is the validation loss at epoch t.

3.6 Performance metrics

To evaluate the dual-stage pneumothorax detection approach, a comprehensive metric system that evaluates the pixel-wise segmentation accuracy and classification reliability is needed. We employ a comprehensive evaluation strategy that considers both spatial and category correctness, as well as the therapeutic significance of different sorts of errors.

3.6.1 Segmentation metrics

For evaluating segmentation performance, we employ a suite of complementary metrics that capture different aspects of spatial accuracy. Let $M_{\text {pred}}$ and $M_{\text {true}}$ represent the predicted and ground truth segmentation masks respectively.

The Intersection over Union (IoU), also known as the Jaccard Index, is calculated as:

$I o U=\frac{\left|M_{\text {pred }} \cap M_{\text {true }}\right|}{\left|M_{\text {pred }} \cup M_{\text {true }}\right|}$            (21)

3.6.2 Classification metrics

The classification performance is evaluated through a comprehensive set of metrics derived from the confusion matrix. Let TP, TN, FP, and FN denote true positives, true negatives, false positives, and false negatives respectively.

For binary classification metrics:

$\begin{gathered}\text { Precision }=\frac{T P}{T P+F P} \\ \text { Recall }=\frac{T P}{T P+F N} \\ F 1=2 \cdot \frac{\text { Precision ⋅ } \text { Recall }}{\text { Precision }+ \text { Recall }}\end{gathered}$          (22)

additional classification metrics include:

Accuracy:

Accuracy $=\frac{T P+T N}{T P+T N+F P+F N}$          (23)

where, TPR is the true positive rate and FPR is the false positive rate.

This study presents a comprehensive dual-stage framework for pneumothorax detection and segmentation in chest radiographs, introducing several novel contributions to the field of automated medical image analysis. The experimental results offer valuable insights into both the technical aspects of deep learning approaches for medical image analysis and their potential clinical implications.

An interesting result was discovered from the comparison of backbone architectures: InceptionResNetV2 had the lowest Intersection over Union score (0.2238) despite having the highest overall accuracy (0.9943), indicating the drawbacks of using accuracy as the only evaluation metric in extremely unbalanced segmentation tasks. The difference emphasizes how important it is to employ boundary-sensitive measurements, such as Intersection over Union (IoU), for medical segmentation tasks, especially when it comes to situations like pneumothorax.

All models, including the top-performing Xceptio Xception based architecture, occasionally had difficulty recognizing small or small pneurthoraces, especially those that presented atypically or in cases with poor quality images, according to the qualitative analysis of segmentation results. This restriction highlights the inherent difficulties in interpreting chest X-ray and points to possible areas for development by incorporating attention mechanisms made especially to improve focus on areas with subtle intensity variations.

A significant difference between sensitivity and specificity was identified in the stage of classification by comparing the performance metrics of the XGBoost model before and after hyperparameter changes. By improving the ability to detect positive instances, the optimization procedure raised the recall for pneumothorax cases from 0.64 to 0.76. This is especially useful in emergency situations where failing to detect a pneumothorax could have serious clinical repercussions. However, this improvement came with a slight decrease in precision (from 0.85 to 0.81), reflecting the inherent challenge in balancing false positives and false negatives in clinical diagnostic systems.

The dual-stage approach combining segmentation with classification demonstrates several advantages over single-stage systems. Using segmentation to first localize pneumothorax areas, and then classification using both segmentation masks and original image attributes, the framework efficiently integrates spatial and contextual information. This integrated in our approach reduced false negatives if we compared to image-only classification methods reported in prior studies such as Kaur et al. [17], who stated that their respective recall ratings for pneumothorax detection were 0.71 and 0.73.