Duplicate from szabopeter/roi-segresnet-2d

Co-authored-by: Werner Peter <szabopeter@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,115 @@
+---
+license: apache-2.0
+tags:
+- medical-imaging
+- segmentation
+- image-segmentation
+- ct
+- lung
+- monai
+- pytorch
+library_name: pytorch
+pipeline_tag: image-segmentation
+---
+
+# Lung ROI Segmentation — SegResNet (2D)
+
+Per-slice 2D lung foreground (thoracic parenchyma) segmentation for
+chest CT. Used as the ROI stage of a two-stage coarse-to-fine
+pulmonary-nodule segmentation pipeline: its outputs are stacked into a
+3D lung bounding box that the downstream nodule model crops to.
+
+## Model details
+
+- **Architecture**: [MONAI SegResNet](https://docs.monai.io/en/stable/networks.html#segresnet), 2D residual U-Net
+- **Trainable parameters**: 6,904,081
+- **Input**: `(1, 256, 256)` axial CT slice, intensity-normalised to `[0, 1]`
+- **Output**: `(1, 256, 256)` sigmoid map; foreground = lung tissue
+- **Framework**: PyTorch + MONAI
+
+## Data
+
+Trained on the `unified` split (patient-grouped, dataset-stratified)
+of a unified corpus assembled from three public sources:
+
+- NLST
+- NSCLC-Radiomics
+- LIDC-IDRI
+
+Slices per split: **365 014 train / 64 097 val / — test held out**.
+
+## Validation metrics
+
+Evaluated on 64 097 val slices (≈ 4.2 × 10⁹ pixels), micro-averaged at
+0.5 threshold on sigmoid output.
+
+| Metric      | Value  |
+|-------------|-------:|
+| mIoU        | 0.9803 |
+| Accuracy    | 0.9953 |
+| Precision   | 0.9797 |
+| Recall      | 0.9857 |
+| Dice (F1)   | 0.9827 |
+
+## How to load & run inference
+
+```python
+import yaml, torch
+from monai.networks.nets import SegResNet
+
+cfg = yaml.safe_load(open("config.yaml"))["model"]
+model = SegResNet(
+    spatial_dims = cfg["spatial_dims"],
+    in_channels  = cfg["in_channels"],
+    out_channels = cfg["out_channels"],
+    init_filters = cfg["init_filters"],
+    blocks_down  = tuple(cfg["blocks_down"]),
+    blocks_up    = tuple(cfg["blocks_up"]),
+    dropout_prob = cfg["dropout_prob"],
+)
+state = torch.load("model.pth", map_location="cpu", weights_only=True)
+model.load_state_dict(state)
+model.eval()
+
+with torch.no_grad():
+    x = torch.randn(1, 1, 256, 256)         # (B, C, H, W) — replace with your CT slice
+    prob = torch.sigmoid(model(x))
+    lung_mask = (prob > 0.5).to(torch.uint8)
+```
+
+## Training recipe
+
+- **Loss**: `DiceLoss(sigmoid=True, squared_pred=True)`
+- **Optimizer**: Adam (lr = 1e-3, weight decay = 1e-5)
+- **Scheduler**: CosineAnnealingLR (T_max = 100, η_min = 1e-6)
+- **Batch size**: 16
+- **Samples/epoch**: 20 000 (random subsample of ~365 k train slices)
+- **Epochs budget**: 100  |  **best checkpoint at epoch 7 / 100**
+- **Mixed precision**: bf16
+- **Seed**: 42
+- **Hardware**: 1 × NVIDIA H100 94 GB
+
+Full config is included in this repo as `config.yaml`.
+
+## Reproducing training
+
+Training code lives in an accompanying reproduction demo (published
+separately). Once available, reproduce with:
+
+```bash
+export DATA_ROOT=/path/to/unified          # dir containing ct_2d/ and roi_sem_seg_2d/
+python train.py --config config.yaml
+```
+
+## License & intended use
+
+Model weights released under Apache 2.0. Training data was public but
+covered by dataset-specific terms (NLST, NSCLC-Radiomics, LIDC-IDRI) —
+users must comply with those separately when using the model on
+comparable data.
+
+**Not a medical device.** Not intended for clinical use. Research only.
+
+## Citation
+
+Paper in preparation.

config.yaml

@@ -0,0 +1,45 @@
+# ROI — 2D lung foreground segmentation.
+# Reproduces the bundled checkpoint at checkpoints/roi/best.pth
+# (val Dice = 0.96 on the unified split).
+
+task: roi
+
+data:
+  data_root:   ${DATA_ROOT}           # contains ct_2d/ and roi_sem_seg_2d/
+  split_json:  data/splits/unified.json
+
+preprocessing:
+  target_size:  [256, 256]
+
+model:
+  name:         segresnet
+  spatial_dims: 2
+  in_channels:  1
+  out_channels: 1                     # sigmoid
+  init_filters: 16
+  blocks_down:  [1, 2, 2, 4, 4]
+  blocks_up:    [1, 1, 1, 1]
+  dropout_prob: 0.1
+
+training:
+  epochs:            100
+  batch_size:        16
+  samples_per_epoch: 20000            # subsample ~1/24 of ~475k slices
+  lr:                1.0e-3
+  weight_decay:      1.0e-5
+  val_interval:      1
+  device:            cuda
+  amp:               true
+  amp_dtype:         bfloat16
+  seed:              42
+
+scheduler:
+  T_max:   100
+  eta_min: 1.0e-6
+
+checkpoints:
+  dir:        checkpoints/roi
+  save_every: 10
+
+logging:
+  dir:        runs/roi

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0edd572abb3a863f1bf6d71c62da285d7a92998189f0baf98a668660e65d7e3
+size 27732822