vasandesh/pneumothorax-unet-efficientnet-b4

Pneumothorax Detection & Segmentation

U-Net with EfficientNet-B4 backbone for pneumothorax detection and segmentation on chest X-rays.

Model Details

• Architecture: U-Net with ImageNet-pretrained EfficientNet-B4 encoder
• Input: 512×512 chest X-ray images (3-channel)
• Output: Segmentation mask (512×512) + Classification logit

Training

• Dataset: SIIM-ACR Pneumothorax Segmentation (2,379 positive cases)
• Optimizer: AdamW (lr=1e-4)
• Loss: Focal + Dice (segmentation) + BCE (classification)
• Early stopping: After 7 epochs of no improvement
• Best epoch: 10 / 20 total

Performance (Test Set, n=1,602)

• Classification AUC: 0.9464
• Sensitivity: 81.6%, Specificity: 92.1%
• Segmentation Dice: 0.5142 (lesion-positive cases)

Operating Points

Threshold	Sensitivity	Specificity
0.3	85.6%	90.8%
0.5	81.6%	92.1%
0.68	80.2%	93.2%

Limitations

• Single-dataset training (SIIM-ACR only)
• Sensitivity 2.7 pp below GE predicate at matched specificity
• Segmentation quality (Dice 0.51) indicates room for improvement
• Not validated for clinical use without independent verification

Next Steps

• Multi-dataset training on MIMIC-CXR for better generalization
• Threshold tuning for deployment

Citation

Dataset: SIIM-ACR Pneumothorax Segmentation Challenge (Kaggle 2019)

License

MIT - See accompanying repository for details

Disclaimer

Research model only. Not intended for clinical use without regulatory clearance.