sindri101/chest-drain-predictor

Chest Drain Predictor

Binary classifier that detects chest drain (thoracostomy tube) presence on chest X-rays.

Built as an auxiliary confounder predictor for downstream pneumothorax classification tasks, to mitigate shortcut learning from chest drains.

Performance

Metric	Validation	Test
Sensitivity (Recall)	1.0000	1.0000
Specificity	0.9452	0.9521
AUC	0.9997	0.9987
F1	0.9676	0.9643
Precision	0.9373	0.9310

Usage

Quick start

# pip install torch timm torchvision huggingface_hub

from huggingface_hub import snapshot_download
import sys

# Download model (cached after first call)
model_dir = snapshot_download("sindri101/chest-drain-predictor")
sys.path.insert(0, model_dir)

from chest_drain_detector import load_model

detector = load_model()

# Single image
result = detector("path/to/chest_xray.png")
print(result)
# {'prediction': 'chest_drain_present', 'probability': 0.9993, 'label': 1}

# Batch prediction
results = detector.predict_batch(["img1.png", "img2.png", "img3.png"])

From a local directory

sys.path.insert(0, "/path/to/chest_drain_predictor")
from chest_drain_detector import ChestDrainDetector

detector = ChestDrainDetector.from_pretrained("/path/to/chest_drain_detector")
result = detector("path/to/chest_xray.png")

CLI

python chest_drain_detector.py path/to/chest_xray.png

Model Details

• Architecture: ConvNeXt-Tiny (pretrained on ImageNet, fine-tuned)
• Input: Chest X-ray image (PNG/JPEG), any resolution (resized to 512x512 internally)
• Output: Binary prediction (chest drain present / absent) with probability score
• Decision threshold: 0.05 (optimized for high recall)
• Framework: PyTorch + timm

Training Details

• Dataset: NIH ChestX-ray14 with NEATX chest drain annotations (3,543 images, ~46% positive)
• Split: Patient-level stratified 70/15/15
• Loss: Focal loss (alpha=0.6, gamma=2.0)
• Optimizer: AdamW (lr=3e-5, weight_decay=1e-4)
• Training time: 30 minutes on RTX 3090

Dependencies

torch
timm
torchvision
Pillow

Intended Use

This model is designed as a confounder predictor, not a diagnostic tool. Its primary purpose is to flag chest X-rays containing chest drains so that downstream models (e.g., pneumothorax classifiers) can account for this potential shortcut feature.

Limitations

• Trained on NIH ChestX-ray14 only; may not generalize perfectly to other CXR datasets with different acquisition protocols
• Annotations are from non-expert annotators (validated against expert labels with moderate-to-almost-perfect agreement)
• Not intended for clinical diagnostic use