Add portable PE segmentation model

.gitattributes

@@ -1,35 +1 @@
-*.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
+pytorch_model.bin filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,140 @@
+---
+library_name: transformers
+pipeline_tag: image-segmentation
+tags:
+  - ct-pulmonary-angiography
+  - medical-imaging
+  - ct
+  - pulmonary-embolism
+  - segmentation
+  - nnunet
+  - pytorch
+license: apache-2.0
+datasets:
+  - mazurowski-lab/PulmonaryEmbolismSegmentation
+paperswithcode_id: pulmonary-embolism-segmentation
+model-index:
+  - name: Pulmonary Embolism Segmentation
+    results:
+      - task:
+          type: image-segmentation
+          name: 3D medical image segmentation
+        dataset:
+          name: Pulmonary Embolism Segmentation
+          type: mazurowski-lab/PulmonaryEmbolismSegmentation
+        metrics:
+          - type: dice
+            name: Mean Dice
+            value: 0.5357
+---
+
+# Pulmonary Embolism Segmentation
+
+This repository contains a portable Hugging Face-compatible version of the
+Pulmonary Embolism Segmentation model from Mazurowski Lab.
+
+The model is a 3D Residual Encoder U-Net trained with nnU-Net v2. The runtime
+model code is implemented directly in PyTorch, so inference does not require
+`nnunetv2` or `dynamic-network-architectures`.
+
+- GitHub: https://github.com/mazurowski-lab/PulmonaryEmbolismSegmentation
+- Paper: https://link.springer.com/article/10.1007/s10278-026-01958-4
+- Contact: Yixin Zhang, yz696@duke.edu
+
+## Model
+
+- Architecture: 3D ResidualEncoderUNet
+- Input: single-channel CT volume
+- Output: 2 logits, background and pulmonary embolism
+- Training patch size: `[224, 320, 320]`
+- Plan spacing: `[1.0, 0.7373045682907104, 0.7373045682907104]`
+- Published checkpoint: `fold_all`
+
+## Usage
+
+```python
+from pulmonary_embolism_segmentation import PulmonaryEmbolismSegmentationModel
+
+model = PulmonaryEmbolismSegmentationModel.from_pretrained(
+    "PATH_OR_REPO_ID",
+    trust_remote_code=True,
+)
+```
+
+For full DICOM inference, use the included helper code from the project
+repository:
+
+```powershell
+python scripts/run_inference.py `
+  --model-dir hf_pe_segmentation_fold_all `
+  --input sample_data/02GE/dicom `
+  --output outputs/02GE_segmentation_nnunet_preprocess.npz `
+  --tile-size 128,256,256
+```
+
+The inference helper follows the nnU-Net v2 preprocessing order:
+
+1. read image and spacing
+2. crop nonzero region
+3. CT clip and normalize
+4. resample image to plan spacing
+5. sliding-window prediction
+6. resample logits back to cropped source grid
+7. argmax
+8. insert crop back into the original image shape
+
+## Dependencies
+
+Core model loading:
+
+- `torch`
+- `transformers`
+- `numpy`
+- `scipy`
+- `scikit-image`
+
+DICOM/NIfTI helpers:
+
+- `SimpleITK`
+- `nibabel`
+
+## Validation
+
+A 40-case validation pass was run on the pixel-level annotated dataset used for
+testing portability. Dice scores are computed after resampling predictions back
+to the original DICOM grid.
+
+| Group | n | Mean Dice | Median Dice | Min | Max | Dice >= 0.5 | Dice >= 0.7 |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| GE | 20 | 0.5934 | 0.6918 | 0.0016 | 0.8512 | 15 | 10 |
+| Toshiba | 20 | 0.4780 | 0.7103 | 0.0000 | 0.8020 | 11 | 11 |
+| All | 40 | 0.5357 | 0.7035 | 0.0000 | 0.8512 | 26 | 21 |
+
+Low-scoring cases should be reviewed for possible annotation, DICOM series, or
+orientation/metadata issues before drawing model-performance conclusions.
+
+Additional input-axis permutation checks were run on low-performing cases
+(`04TS`, `20GE`, `01TS`, `17TS`). The original `(z, y, x)` input order gave the
+best Dice in each tested case, suggesting these failures are not explained by a
+simple array transpose mismatch.
+
+## Intended Use
+
+This model is intended for research use in pulmonary embolism segmentation from
+CT pulmonary angiography. It is not a medical device and should not be used for
+clinical decision-making without appropriate validation.
+
+## Citation
+
+If you use this model, please cite the associated paper:
+
+```bibtex
+@article{pulmonary_embolism_segmentation_2026,
+  title = {Pulmonary Embolism Segmentation},
+  author = {Zhang, Yixin},
+  journal = {Journal of Imaging Informatics in Medicine},
+  year = {2026},
+  doi = {10.1007/s10278-026-01958-4},
+  url = {https://link.springer.com/article/10.1007/s10278-026-01958-4}
+}
+```

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": ["PulmonaryEmbolismSegmentationModel"],
+  "auto_map": {
+    "AutoConfig": "configuration_pe_segmentation.PulmonaryEmbolismSegmentationConfig",
+    "AutoModel": "modeling_pe_segmentation.PulmonaryEmbolismSegmentationModel"
+  },
+  "model_type": "pulmonary-embolism-segmentation",
+  "input_channels": 1,
+  "num_labels": 2,
+  "patch_size": [224, 320, 320],
+  "spacing": [1.0, 0.7373045682907104, 0.7373045682907104],
+  "features_per_stage": [32, 64, 128, 256, 320, 320, 320],
+  "kernel_sizes": [[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
+  "strides": [[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [1, 2, 2]],
+  "n_blocks_per_stage": [1, 3, 4, 6, 6, 6, 6],
+  "n_conv_per_stage_decoder": [1, 1, 1, 1, 1, 1],
+  "conv_bias": true,
+  "norm_eps": 1e-05,
+  "norm_affine": true,
+  "leaky_relu_inplace": true,
+  "deep_supervision": false,
+  "label2id": {
+    "background": 0,
+    "pulmonary_embolism": 1
+  },
+  "id2label": {
+    "0": "background",
+    "1": "pulmonary_embolism"
+  },
+  "ct_clip_min": -195.0,
+  "ct_clip_max": 305.0,
+  "ct_mean": 37.060203552246094,
+  "ct_std": 92.34374237060547
+}

configuration_pe_segmentation.py

@@ -0,0 +1,61 @@
+from transformers import PretrainedConfig
+
+
+class PulmonaryEmbolismSegmentationConfig(PretrainedConfig):
+    model_type = "pulmonary-embolism-segmentation"
+
+    def __init__(
+        self,
+        input_channels=1,
+        num_labels=2,
+        patch_size=None,
+        spacing=None,
+        features_per_stage=None,
+        kernel_sizes=None,
+        strides=None,
+        n_blocks_per_stage=None,
+        n_conv_per_stage_decoder=None,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        leaky_relu_inplace=True,
+        deep_supervision=False,
+        label2id=None,
+        id2label=None,
+        ct_clip_min=-195.0,
+        ct_clip_max=305.0,
+        ct_mean=37.060203552246094,
+        ct_std=92.34374237060547,
+        **kwargs,
+    ):
+        label2id = label2id or {"background": 0, "pulmonary_embolism": 1}
+        id2label = id2label or {str(v): k for k, v in label2id.items()}
+        super().__init__(label2id=label2id, id2label=id2label, **kwargs)
+
+        self.input_channels = input_channels
+        self.num_labels = num_labels
+        self.patch_size = patch_size or [224, 320, 320]
+        self.spacing = spacing or [1.0, 0.7373045682907104, 0.7373045682907104]
+        self.features_per_stage = features_per_stage or [32, 64, 128, 256, 320, 320, 320]
+        self.kernel_sizes = kernel_sizes or [[3, 3, 3]] * 7
+        self.strides = strides or [
+            [1, 1, 1],
+            [2, 2, 2],
+            [2, 2, 2],
+            [2, 2, 2],
+            [2, 2, 2],
+            [2, 2, 2],
+            [1, 2, 2],
+        ]
+        self.n_blocks_per_stage = n_blocks_per_stage or [1, 3, 4, 6, 6, 6, 6]
+        self.n_conv_per_stage_decoder = n_conv_per_stage_decoder or [1, 1, 1, 1, 1, 1]
+        self.conv_bias = conv_bias
+        self.norm_eps = norm_eps
+        self.norm_affine = norm_affine
+        self.leaky_relu_inplace = leaky_relu_inplace
+        self.deep_supervision = deep_supervision
+
+        self.ct_clip_min = ct_clip_min
+        self.ct_clip_max = ct_clip_max
+        self.ct_mean = ct_mean
+        self.ct_std = ct_std

local_architecture.py

@@ -0,0 +1,344 @@
+from typing import Sequence
+
+import torch
+from torch import nn
+
+
+def _as_tuple(value):
+    if isinstance(value, int):
+        return (value, value, value)
+    return tuple(value)
+
+
+class ConvDropoutNormReLU(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        output_channels,
+        kernel_size,
+        stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        with_nonlin=True,
+    ):
+        super().__init__()
+        kernel_size = _as_tuple(kernel_size)
+        stride = _as_tuple(stride)
+        padding = tuple((k - 1) // 2 for k in kernel_size)
+        self.conv = nn.Conv3d(input_channels, output_channels, kernel_size, stride, padding, bias=conv_bias)
+        self.norm = nn.InstanceNorm3d(output_channels, eps=norm_eps, affine=norm_affine)
+        modules = [self.conv, self.norm]
+        if with_nonlin:
+            self.nonlin = nn.LeakyReLU(negative_slope=0.01, inplace=True)
+            modules.append(self.nonlin)
+        self.all_modules = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.all_modules(x)
+
+
+class StackedConvBlocks(nn.Module):
+    def __init__(
+        self,
+        num_convs,
+        input_channels,
+        output_channels,
+        kernel_size,
+        initial_stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        if not isinstance(output_channels, (tuple, list)):
+            output_channels = [output_channels] * num_convs
+        blocks = [
+            ConvDropoutNormReLU(
+                input_channels,
+                output_channels[0],
+                kernel_size,
+                initial_stride,
+                conv_bias=conv_bias,
+                norm_eps=norm_eps,
+                norm_affine=norm_affine,
+                with_nonlin=True,
+            )
+        ]
+        for idx in range(1, num_convs):
+            blocks.append(
+                ConvDropoutNormReLU(
+                    output_channels[idx - 1],
+                    output_channels[idx],
+                    kernel_size,
+                    1,
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                    with_nonlin=True,
+                )
+            )
+        self.convs = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.convs(x)
+
+
+class BasicBlockD(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        output_channels,
+        kernel_size,
+        stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        stride = _as_tuple(stride)
+        self.conv1 = ConvDropoutNormReLU(
+            input_channels,
+            output_channels,
+            kernel_size,
+            stride,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+            with_nonlin=True,
+        )
+        self.conv2 = ConvDropoutNormReLU(
+            output_channels,
+            output_channels,
+            kernel_size,
+            1,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+            with_nonlin=False,
+        )
+        self.nonlin2 = nn.LeakyReLU(negative_slope=0.01, inplace=True)
+
+        has_stride = any(s != 1 for s in stride)
+        requires_projection = input_channels != output_channels
+        if has_stride or requires_projection:
+            ops = []
+            if has_stride:
+                ops.append(nn.AvgPool3d(kernel_size=stride, stride=stride))
+            if requires_projection:
+                ops.append(
+                    ConvDropoutNormReLU(
+                        input_channels,
+                        output_channels,
+                        1,
+                        1,
+                        conv_bias=False,
+                        norm_eps=norm_eps,
+                        norm_affine=norm_affine,
+                        with_nonlin=False,
+                    )
+                )
+            self.skip = nn.Sequential(*ops)
+        else:
+            self.skip = nn.Identity()
+
+    def forward(self, x):
+        return self.nonlin2(self.conv2(self.conv1(x)) + self.skip(x))
+
+
+class StackedResidualBlocks(nn.Module):
+    def __init__(
+        self,
+        n_blocks,
+        input_channels,
+        output_channels,
+        kernel_size,
+        initial_stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        blocks = [
+            BasicBlockD(
+                input_channels,
+                output_channels,
+                kernel_size,
+                initial_stride,
+                conv_bias=conv_bias,
+                norm_eps=norm_eps,
+                norm_affine=norm_affine,
+            )
+        ]
+        for _ in range(1, n_blocks):
+            blocks.append(
+                BasicBlockD(
+                    output_channels,
+                    output_channels,
+                    kernel_size,
+                    1,
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                )
+            )
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.blocks(x)
+
+
+class ResidualEncoder(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        features_per_stage,
+        kernel_sizes,
+        strides,
+        n_blocks_per_stage,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        self.stem = StackedConvBlocks(
+            1,
+            input_channels,
+            features_per_stage[0],
+            kernel_sizes[0],
+            1,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+        )
+        input_channels = features_per_stage[0]
+        stages = []
+        for idx, output_channels in enumerate(features_per_stage):
+            stages.append(
+                StackedResidualBlocks(
+                    n_blocks_per_stage[idx],
+                    input_channels,
+                    output_channels,
+                    kernel_sizes[idx],
+                    strides[idx],
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                )
+            )
+            input_channels = output_channels
+        self.stages = nn.Sequential(*stages)
+
+    def forward(self, x):
+        x = self.stem(x)
+        skips = []
+        for stage in self.stages:
+            x = stage(x)
+            skips.append(x)
+        return skips
+
+
+class UNetDecoder(nn.Module):
+    def __init__(
+        self,
+        features_per_stage: Sequence[int],
+        strides: Sequence[Sequence[int]],
+        num_classes: int,
+        n_conv_per_stage_decoder: Sequence[int],
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        deep_supervision=False,
+    ):
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.encoder = nn.Identity()
+        encoder_channels = list(features_per_stage)
+        decoder_channels = list(features_per_stage[:-1][::-1])
+        bottleneck_channels = features_per_stage[-1]
+        transpose_strides = list(strides[1:][::-1])
+
+        self.stages = nn.ModuleList()
+        self.transpconvs = nn.ModuleList()
+        self.seg_layers = nn.ModuleList()
+
+        input_channels = bottleneck_channels
+        for idx, output_channels in enumerate(decoder_channels):
+            self.transpconvs.append(
+                nn.ConvTranspose3d(
+                    input_channels,
+                    output_channels,
+                    kernel_size=_as_tuple(transpose_strides[idx]),
+                    stride=_as_tuple(transpose_strides[idx]),
+                )
+            )
+            self.stages.append(
+                StackedConvBlocks(
+                    n_conv_per_stage_decoder[idx],
+                    output_channels + encoder_channels[-(idx + 2)],
+                    output_channels,
+                    3,
+                    1,
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                )
+            )
+            self.seg_layers.append(nn.Conv3d(output_channels, num_classes, 1, 1, 0))
+            input_channels = output_channels
+
+    def forward(self, skips):
+        x = skips[-1]
+        seg_outputs = []
+        for idx, stage in enumerate(self.stages):
+            x = self.transpconvs[idx](x)
+            x = torch.cat((x, skips[-(idx + 2)]), dim=1)
+            x = stage(x)
+            if self.deep_supervision:
+                seg_outputs.append(self.seg_layers[idx](x))
+            elif idx == len(self.stages) - 1:
+                seg_outputs.append(self.seg_layers[-1](x))
+        seg_outputs = seg_outputs[::-1]
+        return seg_outputs if self.deep_supervision else seg_outputs[0]
+
+
+class ResidualEncoderUNet(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        features_per_stage,
+        kernel_sizes,
+        strides,
+        n_blocks_per_stage,
+        num_classes,
+        n_conv_per_stage_decoder,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        deep_supervision=False,
+    ):
+        super().__init__()
+        self.encoder = ResidualEncoder(
+            input_channels,
+            features_per_stage,
+            kernel_sizes,
+            strides,
+            n_blocks_per_stage,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+        )
+        self.decoder = UNetDecoder(
+            features_per_stage,
+            strides,
+            num_classes,
+            n_conv_per_stage_decoder,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+            deep_supervision=deep_supervision,
+        )
+        self.decoder.encoder = self.encoder
+
+    def forward(self, x):
+        return self.decoder(self.encoder(x))

modeling_pe_segmentation.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+try:
+    from .configuration_pe_segmentation import PulmonaryEmbolismSegmentationConfig
+except ImportError:
+    from configuration_pe_segmentation import PulmonaryEmbolismSegmentationConfig
+
+
+@dataclass
+class SegmentationModelOutput(ModelOutput):
+    loss: Optional[torch.Tensor] = None
+    logits: torch.Tensor = None
+    deep_supervision_logits: Optional[Tuple[torch.Tensor, ...]] = None
+
+
+class PulmonaryEmbolismSegmentationModel(PreTrainedModel):
+    config_class = PulmonaryEmbolismSegmentationConfig
+    base_model_prefix = "segmentation_model"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = False
+    _tied_weights_keys = []
+    all_tied_weights_keys = {}
+    _keys_to_ignore_on_load_missing = [
+        r"segmentation_model\..*\.all_modules\..*",
+        r"segmentation_model\.decoder\.encoder\..*",
+    ]
+
+    def __init__(self, config: PulmonaryEmbolismSegmentationConfig):
+        super().__init__(config)
+        self.segmentation_model = self._build_network(config)
+
+    @staticmethod
+    def _build_network(config: PulmonaryEmbolismSegmentationConfig) -> nn.Module:
+        try:
+            from .local_architecture import ResidualEncoderUNet
+        except ImportError:
+            from pulmonary_embolism_segmentation.local_architecture import ResidualEncoderUNet
+
+        return ResidualEncoderUNet(
+            input_channels=config.input_channels,
+            features_per_stage=config.features_per_stage,
+            kernel_sizes=config.kernel_sizes,
+            strides=config.strides,
+            n_blocks_per_stage=config.n_blocks_per_stage,
+            num_classes=config.num_labels,
+            n_conv_per_stage_decoder=config.n_conv_per_stage_decoder,
+            conv_bias=config.conv_bias,
+            norm_eps=config.norm_eps,
+            norm_affine=config.norm_affine,
+            deep_supervision=config.deep_supervision,
+        )
+
+    def forward(self, pixel_values: torch.Tensor, labels: Optional[torch.Tensor] = None):
+        outputs = self.segmentation_model(pixel_values)
+        if isinstance(outputs, (tuple, list)):
+            logits = outputs[0]
+            deep_supervision_logits = tuple(outputs[1:])
+        else:
+            logits = outputs
+            deep_supervision_logits = None
+
+        loss = None
+        if labels is not None:
+            loss = nn.functional.cross_entropy(logits, labels.long())
+
+        return SegmentationModelOutput(
+            loss=loss,
+            logits=logits,
+            deep_supervision_logits=deep_supervision_logits,
+        )

pulmonary_embolism_segmentation/__init__.py

@@ -0,0 +1,11 @@
+from .configuration_pe_segmentation import PulmonaryEmbolismSegmentationConfig
+from .inference import SlidingWindowInferenceConfig, predict_volume, predict_volume_resampled
+from .modeling_pe_segmentation import PulmonaryEmbolismSegmentationModel
+
+__all__ = [
+    "PulmonaryEmbolismSegmentationConfig",
+    "PulmonaryEmbolismSegmentationModel",
+    "SlidingWindowInferenceConfig",
+    "predict_volume",
+    "predict_volume_resampled",
+]

pulmonary_embolism_segmentation/configuration_pe_segmentation.py

@@ -0,0 +1,61 @@
+from transformers import PretrainedConfig
+
+
+class PulmonaryEmbolismSegmentationConfig(PretrainedConfig):
+    model_type = "pulmonary-embolism-segmentation"
+
+    def __init__(
+        self,
+        input_channels=1,
+        num_labels=2,
+        patch_size=None,
+        spacing=None,
+        features_per_stage=None,
+        kernel_sizes=None,
+        strides=None,
+        n_blocks_per_stage=None,
+        n_conv_per_stage_decoder=None,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        leaky_relu_inplace=True,
+        deep_supervision=False,
+        label2id=None,
+        id2label=None,
+        ct_clip_min=-195.0,
+        ct_clip_max=305.0,
+        ct_mean=37.060203552246094,
+        ct_std=92.34374237060547,
+        **kwargs,
+    ):
+        label2id = label2id or {"background": 0, "pulmonary_embolism": 1}
+        id2label = id2label or {str(v): k for k, v in label2id.items()}
+        super().__init__(label2id=label2id, id2label=id2label, **kwargs)
+
+        self.input_channels = input_channels
+        self.num_labels = num_labels
+        self.patch_size = patch_size or [224, 320, 320]
+        self.spacing = spacing or [1.0, 0.7373045682907104, 0.7373045682907104]
+        self.features_per_stage = features_per_stage or [32, 64, 128, 256, 320, 320, 320]
+        self.kernel_sizes = kernel_sizes or [[3, 3, 3]] * 7
+        self.strides = strides or [
+            [1, 1, 1],
+            [2, 2, 2],
+            [2, 2, 2],
+            [2, 2, 2],
+            [2, 2, 2],
+            [2, 2, 2],
+            [1, 2, 2],
+        ]
+        self.n_blocks_per_stage = n_blocks_per_stage or [1, 3, 4, 6, 6, 6, 6]
+        self.n_conv_per_stage_decoder = n_conv_per_stage_decoder or [1, 1, 1, 1, 1, 1]
+        self.conv_bias = conv_bias
+        self.norm_eps = norm_eps
+        self.norm_affine = norm_affine
+        self.leaky_relu_inplace = leaky_relu_inplace
+        self.deep_supervision = deep_supervision
+
+        self.ct_clip_min = ct_clip_min
+        self.ct_clip_max = ct_clip_max
+        self.ct_mean = ct_mean
+        self.ct_std = ct_std

pulmonary_embolism_segmentation/inference.py

@@ -0,0 +1,362 @@
+from dataclasses import dataclass
+from itertools import product
+from math import ceil
+from typing import Iterable, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy.ndimage import binary_fill_holes, gaussian_filter, map_coordinates
+from skimage.transform import resize
+
+
+ANISO_THRESHOLD = 3.0
+
+
+def normalize_ct(volume: np.ndarray, config) -> np.ndarray:
+    volume = volume.astype(np.float32, copy=False)
+    volume = np.clip(volume, config.ct_clip_min, config.ct_clip_max)
+    return (volume - config.ct_mean) / config.ct_std
+
+
+def compute_new_shape(old_shape: Sequence[int], old_spacing: Sequence[float], new_spacing: Sequence[float]) -> Tuple[int, ...]:
+    return tuple(int(round(spacing / target_spacing * size)) for size, spacing, target_spacing in zip(old_shape, old_spacing, new_spacing))
+
+
+def get_lowres_axis(spacing: Sequence[float]) -> Optional[int]:
+    spacing = np.asarray(spacing)
+    axes = np.where(np.max(spacing) / spacing == 1)[0]
+    return int(axes[0]) if len(axes) == 1 else None
+
+
+def determine_separate_z_axis(current_spacing: Sequence[float], new_spacing: Sequence[float]) -> Optional[int]:
+    current_spacing = np.asarray(current_spacing)
+    new_spacing = np.asarray(new_spacing)
+    if np.max(current_spacing) / np.min(current_spacing) > ANISO_THRESHOLD:
+        return get_lowres_axis(current_spacing)
+    if np.max(new_spacing) / np.min(new_spacing) > ANISO_THRESHOLD:
+        return get_lowres_axis(new_spacing)
+    return None
+
+
+def create_nonzero_mask(data: np.ndarray) -> np.ndarray:
+    mask = data[0] != 0
+    for channel in range(1, data.shape[0]):
+        mask |= data[channel] != 0
+    return binary_fill_holes(mask)
+
+
+def crop_to_nonzero(data: np.ndarray) -> Tuple[np.ndarray, Tuple[Tuple[int, int], ...]]:
+    mask = create_nonzero_mask(data)
+    coords = np.array(np.where(mask))
+    if coords.size == 0:
+        bbox = tuple((0, dim) for dim in data.shape[1:])
+    else:
+        mins = coords.min(axis=1)
+        maxs = coords.max(axis=1) + 1
+        bbox = tuple((int(lo), int(hi)) for lo, hi in zip(mins, maxs))
+    slicer = (slice(None),) + tuple(slice(lo, hi) for lo, hi in bbox)
+    return data[slicer], bbox
+
+
+def insert_crop(segmentation: np.ndarray, bbox: Sequence[Tuple[int, int]], original_shape: Sequence[int]) -> np.ndarray:
+    result = np.zeros(original_shape, dtype=segmentation.dtype)
+    slicer = tuple(slice(lo, hi) for lo, hi in bbox)
+    result[slicer] = segmentation
+    return result
+
+
+def _resize_segmentation(segmentation: np.ndarray, new_shape: Sequence[int], order: int) -> np.ndarray:
+    if order == 0:
+        return resize(segmentation, new_shape, order=0, mode="edge", anti_aliasing=False, preserve_range=True)
+    result = np.zeros(new_shape, dtype=segmentation.dtype)
+    for label in np.unique(segmentation):
+        resized = resize(
+            (segmentation == label).astype(float),
+            new_shape,
+            order=order,
+            mode="edge",
+            anti_aliasing=False,
+            preserve_range=True,
+        )
+        result[resized >= 0.5] = label
+    return result
+
+
+def _resize_channel(channel: np.ndarray, new_shape: Sequence[int], is_seg: bool, order: int) -> np.ndarray:
+    if is_seg:
+        return _resize_segmentation(channel, new_shape, order)
+    return resize(channel, new_shape, order=order, mode="edge", anti_aliasing=False, preserve_range=True)
+
+
+def resample_nnunet(
+    data: np.ndarray,
+    new_shape: Sequence[int],
+    current_spacing: Sequence[float],
+    new_spacing: Sequence[float],
+    is_seg: bool = False,
+    order: int = 3,
+    order_z: int = 0,
+) -> np.ndarray:
+    if data is None:
+        return None
+    if isinstance(data, torch.Tensor):
+        data = data.cpu().numpy()
+    assert data.ndim == 4, "data must be channel-first: (c, z, y, x)"
+
+    old_shape = tuple(data.shape[1:])
+    new_shape = tuple(int(i) for i in new_shape)
+    if old_shape == new_shape:
+        return data
+
+    axis = determine_separate_z_axis(current_spacing, new_spacing)
+    output = np.zeros((data.shape[0], *new_shape), dtype=data.dtype)
+    data_float = data.astype(float, copy=False)
+
+    if axis is None:
+        for c in range(data.shape[0]):
+            output[c] = _resize_channel(data_float[c], new_shape, is_seg, order)
+        return output
+
+    shape = np.array(old_shape)
+    new_shape_array = np.array(new_shape)
+    if axis == 0:
+        plane_shape = new_shape_array[1:]
+    elif axis == 1:
+        plane_shape = new_shape_array[[0, 2]]
+    else:
+        plane_shape = new_shape_array[:-1]
+
+    for c in range(data.shape[0]):
+        intermediate_shape = new_shape_array.copy()
+        intermediate_shape[axis] = shape[axis]
+        intermediate = np.zeros(tuple(intermediate_shape), dtype=float)
+        for idx in range(shape[axis]):
+            if axis == 0:
+                intermediate[idx] = _resize_channel(data_float[c, idx], plane_shape, is_seg, order)
+            elif axis == 1:
+                intermediate[:, idx] = _resize_channel(data_float[c, :, idx], plane_shape, is_seg, order)
+            else:
+                intermediate[:, :, idx] = _resize_channel(data_float[c, :, :, idx], plane_shape, is_seg, order)
+
+        if shape[axis] == new_shape_array[axis]:
+            output[c] = intermediate
+            continue
+
+        rows, cols, depth = new_shape
+        orig_rows, orig_cols, orig_depth = intermediate.shape
+        row_scale = orig_rows / rows
+        col_scale = orig_cols / cols
+        depth_scale = orig_depth / depth
+        map_rows, map_cols, map_depth = np.mgrid[:rows, :cols, :depth]
+        coord_map = np.array(
+            [
+                row_scale * (map_rows + 0.5) - 0.5,
+                col_scale * (map_cols + 0.5) - 0.5,
+                depth_scale * (map_depth + 0.5) - 0.5,
+            ]
+        )
+        if not is_seg or order_z == 0:
+            output[c] = map_coordinates(intermediate, coord_map, order=order_z, mode="nearest")
+        else:
+            for label in np.unique(intermediate):
+                output[c][
+                    np.round(map_coordinates((intermediate == label).astype(float), coord_map, order=order_z, mode="nearest")) > 0.5
+                ] = label
+    return output
+
+
+def compute_steps(image_size: Sequence[int], tile_size: Sequence[int], overlap: float) -> Tuple[Tuple[int, ...], ...]:
+    steps = []
+    for image_dim, tile_dim in zip(image_size, tile_size):
+        if image_dim <= tile_dim:
+            steps.append((0,))
+            continue
+        target_step = max(1, int(tile_dim * (1.0 - overlap)))
+        num_steps = int(ceil((image_dim - tile_dim) / target_step)) + 1
+        actual_step = (image_dim - tile_dim) / max(1, num_steps - 1)
+        steps.append(tuple(int(round(actual_step * i)) for i in range(num_steps)))
+    return tuple(steps)
+
+
+def gaussian_importance_map(tile_size: Sequence[int], sigma_scale: float = 1.0 / 8.0) -> torch.Tensor:
+    tmp = np.zeros(tuple(tile_size), dtype=np.float32)
+    tmp[tuple(size // 2 for size in tile_size)] = 1
+    sigmas = [size * sigma_scale for size in tile_size]
+    weight = gaussian_filter(tmp, sigmas, 0, mode="constant", cval=0)
+    weight /= np.max(weight)
+    tensor = torch.from_numpy(weight.astype(np.float32, copy=False))
+    tensor[tensor == 0] = torch.min(tensor[tensor != 0])
+    return tensor
+
+
+def pad_to_tile_size(volume: torch.Tensor, tile_size: Sequence[int]) -> Tuple[torch.Tensor, Tuple[slice, slice, slice]]:
+    spatial = volume.shape[-3:]
+    pad_after = [max(tile - dim, 0) for dim, tile in zip(spatial, tile_size)]
+    if any(pad_after):
+        volume = F.pad(volume, (0, pad_after[2], 0, pad_after[1], 0, pad_after[0]))
+    crop = tuple(slice(0, dim) for dim in spatial)
+    return volume, crop
+
+
+def iter_tiles(image_size: Sequence[int], tile_size: Sequence[int], overlap: float) -> Iterable[Tuple[slice, slice, slice]]:
+    for z, y, x in product(*compute_steps(image_size, tile_size, overlap)):
+        yield (
+            slice(z, z + tile_size[0]),
+            slice(y, y + tile_size[1]),
+            slice(x, x + tile_size[2]),
+        )
+
+
+def shrink_tile(tile_size: Sequence[int], min_tile_size: Sequence[int]) -> Optional[Tuple[int, int, int]]:
+    divisibility = (32, 64, 64)
+    candidates = []
+    for axis, value in enumerate(tile_size):
+        if value > min_tile_size[axis]:
+            next_tile = list(tile_size)
+            step = divisibility[axis]
+            next_value = max(min_tile_size[axis], int(value * 0.75) // step * step)
+            next_tile[axis] = next_value
+            candidates.append(tuple(next_tile))
+    if not candidates:
+        return None
+    return min(candidates, key=np.prod)
+
+
+def normalize_tile_size(
+    image_size: Sequence[int],
+    requested_tile_size: Sequence[int],
+    divisibility: Sequence[int] = (32, 64, 64),
+) -> Tuple[int, int, int]:
+    tile_size = []
+    for image_dim, requested_dim, divisor in zip(image_size, requested_tile_size, divisibility):
+        if requested_dim % divisor != 0:
+            raise ValueError(
+                f"Tile size {tuple(requested_tile_size)} is incompatible with network strides. "
+                f"Expected divisibility by {tuple(divisibility)} for z,y,x."
+            )
+        if image_dim >= requested_dim:
+            tile_size.append(requested_dim)
+        else:
+            tile_size.append(int(ceil(image_dim / divisor) * divisor))
+    return tuple(tile_size)
+
+
+@dataclass
+class SlidingWindowInferenceConfig:
+    tile_size: Tuple[int, int, int] = (128, 256, 256)
+    min_tile_size: Tuple[int, int, int] = (64, 128, 128)
+    overlap: float = 0.5
+    use_amp: bool = True
+    device: Optional[str] = None
+    empty_cache_between_tiles: bool = False
+
+
+def _predict_with_tile_size(
+    model,
+    volume: torch.Tensor,
+    tile_size: Sequence[int],
+    overlap: float,
+    device: torch.device,
+    use_amp: bool,
+    empty_cache_between_tiles: bool,
+) -> torch.Tensor:
+    volume, crop = pad_to_tile_size(volume, tile_size)
+    image_size = volume.shape[-3:]
+    num_classes = model.config.num_labels
+
+    logits_sum = torch.zeros((num_classes, *image_size), dtype=torch.float32, device="cpu")
+    weight_sum = torch.zeros(image_size, dtype=torch.float32, device="cpu")
+    importance = gaussian_importance_map(tile_size).to(device)
+
+    model.eval()
+    with torch.inference_mode():
+        for tile in iter_tiles(image_size, tile_size, overlap):
+            patch = volume[(slice(None), slice(None), *tile)].to(device, non_blocking=True)
+            amp_enabled = use_amp and device.type == "cuda"
+            with torch.autocast(device_type=device.type, enabled=amp_enabled):
+                logits = model(patch).logits[0]
+            weighted_logits = (logits.float() * importance).cpu()
+            logits_sum[(slice(None), *tile)] += weighted_logits
+            weight_sum[tile] += importance.cpu()
+            del patch, logits, weighted_logits
+            if empty_cache_between_tiles and device.type == "cuda":
+                torch.cuda.empty_cache()
+
+    logits_sum /= torch.clamp(weight_sum.unsqueeze(0), min=1e-6)
+    return logits_sum[(slice(None), *crop)]
+
+
+def predict_logits(model, volume: np.ndarray, inference_config: Optional[SlidingWindowInferenceConfig] = None) -> torch.Tensor:
+    inference_config = inference_config or SlidingWindowInferenceConfig()
+    device = torch.device(inference_config.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    model.to(device)
+
+    tensor = torch.from_numpy(volume.astype(np.float32, copy=False)[None, None])
+    tile_size = normalize_tile_size(tensor.shape[-3:], inference_config.tile_size)
+
+    while True:
+        try:
+            return _predict_with_tile_size(
+                model=model,
+                volume=tensor,
+                tile_size=tile_size,
+                overlap=inference_config.overlap,
+                device=device,
+                use_amp=inference_config.use_amp,
+                empty_cache_between_tiles=inference_config.empty_cache_between_tiles,
+            )
+        except torch.cuda.OutOfMemoryError:
+            if device.type != "cuda":
+                raise
+            torch.cuda.empty_cache()
+            next_tile_size = shrink_tile(tile_size, inference_config.min_tile_size)
+            if next_tile_size is None:
+                raise RuntimeError(
+                    f"CUDA out of memory at minimum tile size {tile_size}. "
+                    "Use CPU inference or lower min_tile_size."
+                )
+            tile_size = next_tile_size
+
+
+def predict_volume(model, volume: np.ndarray, inference_config: Optional[SlidingWindowInferenceConfig] = None) -> np.ndarray:
+    normalized = normalize_ct(volume, model.config)
+    logits = predict_logits(model, normalized, inference_config)
+    return torch.argmax(logits, dim=0).numpy().astype(np.uint8)
+
+
+def predict_volume_resampled(
+    model,
+    volume: np.ndarray,
+    source_spacing_zyx: Sequence[float],
+    inference_config: Optional[SlidingWindowInferenceConfig] = None,
+) -> np.ndarray:
+    target_spacing_zyx = tuple(model.config.spacing)
+    original_shape = volume.shape
+    data = volume.astype(np.float32, copy=False)[None]
+    cropped, bbox = crop_to_nonzero(data)
+    cropped_shape = cropped.shape[1:]
+    normalized = cropped.copy()
+    normalized[0] = normalize_ct(normalized[0], model.config)
+    resampled_shape = compute_new_shape(cropped_shape, source_spacing_zyx, target_spacing_zyx)
+    resampled = resample_nnunet(
+        normalized,
+        new_shape=resampled_shape,
+        current_spacing=source_spacing_zyx,
+        new_spacing=target_spacing_zyx,
+        is_seg=False,
+        order=3,
+        order_z=0,
+    )
+    logits = predict_logits(model, resampled[0], inference_config)
+    logits_cropped = resample_nnunet(
+        logits,
+        new_shape=cropped_shape,
+        current_spacing=target_spacing_zyx,
+        new_spacing=source_spacing_zyx,
+        is_seg=False,
+        order=1,
+        order_z=0,
+    )
+    segmentation_cropped = np.argmax(logits_cropped, axis=0).astype(np.uint8, copy=False)
+    return insert_crop(segmentation_cropped, bbox, original_shape)

pulmonary_embolism_segmentation/local_architecture.py

@@ -0,0 +1,344 @@
+from typing import Sequence
+
+import torch
+from torch import nn
+
+
+def _as_tuple(value):
+    if isinstance(value, int):
+        return (value, value, value)
+    return tuple(value)
+
+
+class ConvDropoutNormReLU(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        output_channels,
+        kernel_size,
+        stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        with_nonlin=True,
+    ):
+        super().__init__()
+        kernel_size = _as_tuple(kernel_size)
+        stride = _as_tuple(stride)
+        padding = tuple((k - 1) // 2 for k in kernel_size)
+        self.conv = nn.Conv3d(input_channels, output_channels, kernel_size, stride, padding, bias=conv_bias)
+        self.norm = nn.InstanceNorm3d(output_channels, eps=norm_eps, affine=norm_affine)
+        modules = [self.conv, self.norm]
+        if with_nonlin:
+            self.nonlin = nn.LeakyReLU(negative_slope=0.01, inplace=True)
+            modules.append(self.nonlin)
+        self.all_modules = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.all_modules(x)
+
+
+class StackedConvBlocks(nn.Module):
+    def __init__(
+        self,
+        num_convs,
+        input_channels,
+        output_channels,
+        kernel_size,
+        initial_stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        if not isinstance(output_channels, (tuple, list)):
+            output_channels = [output_channels] * num_convs
+        blocks = [
+            ConvDropoutNormReLU(
+                input_channels,
+                output_channels[0],
+                kernel_size,
+                initial_stride,
+                conv_bias=conv_bias,
+                norm_eps=norm_eps,
+                norm_affine=norm_affine,
+                with_nonlin=True,
+            )
+        ]
+        for idx in range(1, num_convs):
+            blocks.append(
+                ConvDropoutNormReLU(
+                    output_channels[idx - 1],
+                    output_channels[idx],
+                    kernel_size,
+                    1,
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                    with_nonlin=True,
+                )
+            )
+        self.convs = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.convs(x)
+
+
+class BasicBlockD(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        output_channels,
+        kernel_size,
+        stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        stride = _as_tuple(stride)
+        self.conv1 = ConvDropoutNormReLU(
+            input_channels,
+            output_channels,
+            kernel_size,
+            stride,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+            with_nonlin=True,
+        )
+        self.conv2 = ConvDropoutNormReLU(
+            output_channels,
+            output_channels,
+            kernel_size,
+            1,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+            with_nonlin=False,
+        )
+        self.nonlin2 = nn.LeakyReLU(negative_slope=0.01, inplace=True)
+
+        has_stride = any(s != 1 for s in stride)
+        requires_projection = input_channels != output_channels
+        if has_stride or requires_projection:
+            ops = []
+            if has_stride:
+                ops.append(nn.AvgPool3d(kernel_size=stride, stride=stride))
+            if requires_projection:
+                ops.append(
+                    ConvDropoutNormReLU(
+                        input_channels,
+                        output_channels,
+                        1,
+                        1,
+                        conv_bias=False,
+                        norm_eps=norm_eps,
+                        norm_affine=norm_affine,
+                        with_nonlin=False,
+                    )
+                )
+            self.skip = nn.Sequential(*ops)
+        else:
+            self.skip = nn.Identity()
+
+    def forward(self, x):
+        return self.nonlin2(self.conv2(self.conv1(x)) + self.skip(x))
+
+
+class StackedResidualBlocks(nn.Module):
+    def __init__(
+        self,
+        n_blocks,
+        input_channels,
+        output_channels,
+        kernel_size,
+        initial_stride,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        blocks = [
+            BasicBlockD(
+                input_channels,
+                output_channels,
+                kernel_size,
+                initial_stride,
+                conv_bias=conv_bias,
+                norm_eps=norm_eps,
+                norm_affine=norm_affine,
+            )
+        ]
+        for _ in range(1, n_blocks):
+            blocks.append(
+                BasicBlockD(
+                    output_channels,
+                    output_channels,
+                    kernel_size,
+                    1,
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                )
+            )
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.blocks(x)
+
+
+class ResidualEncoder(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        features_per_stage,
+        kernel_sizes,
+        strides,
+        n_blocks_per_stage,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+    ):
+        super().__init__()
+        self.stem = StackedConvBlocks(
+            1,
+            input_channels,
+            features_per_stage[0],
+            kernel_sizes[0],
+            1,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+        )
+        input_channels = features_per_stage[0]
+        stages = []
+        for idx, output_channels in enumerate(features_per_stage):
+            stages.append(
+                StackedResidualBlocks(
+                    n_blocks_per_stage[idx],
+                    input_channels,
+                    output_channels,
+                    kernel_sizes[idx],
+                    strides[idx],
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                )
+            )
+            input_channels = output_channels
+        self.stages = nn.Sequential(*stages)
+
+    def forward(self, x):
+        x = self.stem(x)
+        skips = []
+        for stage in self.stages:
+            x = stage(x)
+            skips.append(x)
+        return skips
+
+
+class UNetDecoder(nn.Module):
+    def __init__(
+        self,
+        features_per_stage: Sequence[int],
+        strides: Sequence[Sequence[int]],
+        num_classes: int,
+        n_conv_per_stage_decoder: Sequence[int],
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        deep_supervision=False,
+    ):
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.encoder = nn.Identity()
+        encoder_channels = list(features_per_stage)
+        decoder_channels = list(features_per_stage[:-1][::-1])
+        bottleneck_channels = features_per_stage[-1]
+        transpose_strides = list(strides[1:][::-1])
+
+        self.stages = nn.ModuleList()
+        self.transpconvs = nn.ModuleList()
+        self.seg_layers = nn.ModuleList()
+
+        input_channels = bottleneck_channels
+        for idx, output_channels in enumerate(decoder_channels):
+            self.transpconvs.append(
+                nn.ConvTranspose3d(
+                    input_channels,
+                    output_channels,
+                    kernel_size=_as_tuple(transpose_strides[idx]),
+                    stride=_as_tuple(transpose_strides[idx]),
+                )
+            )
+            self.stages.append(
+                StackedConvBlocks(
+                    n_conv_per_stage_decoder[idx],
+                    output_channels + encoder_channels[-(idx + 2)],
+                    output_channels,
+                    3,
+                    1,
+                    conv_bias=conv_bias,
+                    norm_eps=norm_eps,
+                    norm_affine=norm_affine,
+                )
+            )
+            self.seg_layers.append(nn.Conv3d(output_channels, num_classes, 1, 1, 0))
+            input_channels = output_channels
+
+    def forward(self, skips):
+        x = skips[-1]
+        seg_outputs = []
+        for idx, stage in enumerate(self.stages):
+            x = self.transpconvs[idx](x)
+            x = torch.cat((x, skips[-(idx + 2)]), dim=1)
+            x = stage(x)
+            if self.deep_supervision:
+                seg_outputs.append(self.seg_layers[idx](x))
+            elif idx == len(self.stages) - 1:
+                seg_outputs.append(self.seg_layers[-1](x))
+        seg_outputs = seg_outputs[::-1]
+        return seg_outputs if self.deep_supervision else seg_outputs[0]
+
+
+class ResidualEncoderUNet(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        features_per_stage,
+        kernel_sizes,
+        strides,
+        n_blocks_per_stage,
+        num_classes,
+        n_conv_per_stage_decoder,
+        conv_bias=True,
+        norm_eps=1e-5,
+        norm_affine=True,
+        deep_supervision=False,
+    ):
+        super().__init__()
+        self.encoder = ResidualEncoder(
+            input_channels,
+            features_per_stage,
+            kernel_sizes,
+            strides,
+            n_blocks_per_stage,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+        )
+        self.decoder = UNetDecoder(
+            features_per_stage,
+            strides,
+            num_classes,
+            n_conv_per_stage_decoder,
+            conv_bias=conv_bias,
+            norm_eps=norm_eps,
+            norm_affine=norm_affine,
+            deep_supervision=deep_supervision,
+        )
+        self.decoder.encoder = self.encoder
+
+    def forward(self, x):
+        return self.decoder(self.encoder(x))

pulmonary_embolism_segmentation/modeling_pe_segmentation.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_pe_segmentation import PulmonaryEmbolismSegmentationConfig
+
+
+@dataclass
+class SegmentationModelOutput(ModelOutput):
+    loss: Optional[torch.Tensor] = None
+    logits: torch.Tensor = None
+    deep_supervision_logits: Optional[Tuple[torch.Tensor, ...]] = None
+
+
+class PulmonaryEmbolismSegmentationModel(PreTrainedModel):
+    config_class = PulmonaryEmbolismSegmentationConfig
+    base_model_prefix = "segmentation_model"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = False
+    _tied_weights_keys = []
+    all_tied_weights_keys = {}
+    _keys_to_ignore_on_load_missing = [
+        r"segmentation_model\..*\.all_modules\..*",
+        r"segmentation_model\.decoder\.encoder\..*",
+    ]
+
+    def __init__(self, config: PulmonaryEmbolismSegmentationConfig):
+        super().__init__(config)
+        self.segmentation_model = self._build_network(config)
+
+    @staticmethod
+    def _build_network(config: PulmonaryEmbolismSegmentationConfig) -> nn.Module:
+        from .local_architecture import ResidualEncoderUNet
+
+        return ResidualEncoderUNet(
+            input_channels=config.input_channels,
+            features_per_stage=config.features_per_stage,
+            kernel_sizes=config.kernel_sizes,
+            strides=config.strides,
+            n_blocks_per_stage=config.n_blocks_per_stage,
+            num_classes=config.num_labels,
+            n_conv_per_stage_decoder=config.n_conv_per_stage_decoder,
+            conv_bias=config.conv_bias,
+            norm_eps=config.norm_eps,
+            norm_affine=config.norm_affine,
+            deep_supervision=config.deep_supervision,
+        )
+
+    def forward(self, pixel_values: torch.Tensor, labels: Optional[torch.Tensor] = None):
+        outputs = self.segmentation_model(pixel_values)
+        if isinstance(outputs, (tuple, list)):
+            logits = outputs[0]
+            deep_supervision_logits = tuple(outputs[1:])
+        else:
+            logits = outputs
+            deep_supervision_logits = None
+
+        loss = None
+        if labels is not None:
+            loss = nn.functional.cross_entropy(logits, labels.long())
+
+        return SegmentationModelOutput(
+            loss=loss,
+            logits=logits,
+            deep_supervision_logits=deep_supervision_logits,
+        )

pytorch_model.bin

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