You need to agree to share your Hugging Face information to access this model

This repository is publicly accessible, but you have to accept the conditions to access its files and content.

⚠️ WARNING: The previous model version has been deleted. This repository contains the updated version.

MedNeXt-B 3D Breast MRI Model & BraTS 2023 5-Fold Baseline

A 3D convolutional neural network architecture based on MedNeXt-B (v1) designed for volumetric breast MRI analysis. The backbone was pretrained on the multi-modal brain tumor segmentation benchmark (BraTS 2023 GLI 5-fold cross-validation) and fine-tuned for breast tumour segmentation using a private breast MRI dataset. The released model also includes representation and classification components.

This repository provides both:

Breast MRI Model: Exported PyTorch program (model.pt2), full model weights (model.safetensors), and architecture configuration (config.json).
Pretraining Baseline Backbone: The complete 5 cross-validation folds from the BraTS 2023 GLI segmentation baseline (fold_0 to fold_4), providing the pretraining foundation.

Files

1. Breast MRI Model

File	Description
`model.pt2`	Exported PyTorch program (`torch.export`): compiled model graph and weights in a single file; runs self-contained without external model definition code.
`model.safetensors`	Weights of the complete model (encoder + auxiliary mask heads + classification head); loaded with the Python model code below.
`config.json`	Architecture parameters, input dimensions, and model configuration.

model.safetensors contains:

Weights	Component
`encoder.*`	MedNeXt-B encoder (6.87 M parameters): initialized from BraTS 2023 pretraining and adapted on volumetric breast MRI.
`aux_s8.`, `aux_s16.`	Auxiliary volumetric representation heads on stride-8 and stride-16 feature maps.
`proj.`, `head.`	Projection and classification heads.

2. BraTS 2023 Glioma Baseline Folds (Pretraining Foundation)

Files	Description
`fold_{0..4}/model.safetensors`	Network weights per cross-validation fold (weights-only, no optimizer state).
`fold_{0..4}/metadata.json`	Trainer configuration, best EMA Dice score, allowed mirroring axes.
`fold_{0..4}/plans.json`	nnU-Net architecture plan (patch size 128³, BS 2, MedNeXt-B).

Usage

Feature Extractor via Exported Program (`model.pt2`)

The exported program runs directly in PyTorch without requiring custom architecture classes:

import torch
from huggingface_hub import hf_hub_download

# Download from Hugging Face Hub (or provide local file path)
model_path = hf_hub_download("David-MRI-ORG/mednext-3dconvnext-baseline-breast", "model.pt2")
fx = torch.export.load(model_path).module().to("cuda")   # .to("cpu") works too

# Extract feature tokens and token weights
tokens, token_weights = fx(crops, valid, tumour, voi)

Input format per subject (float32, fixed shapes):

Input	Shape	Description
`crops`	`[6, 64, 96, 96]`	Aligned MRI volumes
`valid`	`[6, 64, 96, 96]`	Binary mask (`1` inside scanned field of view, `0` outside)
`tumour`, `voi`	`[2, 64, 96, 96]`	Region masks for lesion and volume-of-interest

Outputs:

Output	Shape	Description
`tokens`	`[6, 144, 768]`	4 × 6 × 6 spatial grid feature tokens (bottleneck 512 + pooled stride-8 map 256)
`token_weights`	`[6, 144]`	Regional spatial weighting per token

⚠️ .pt2 files contain a compiled program graph: load them only from trusted sources.

Complete Model from Weights (`model.safetensors` + Code)

Save the Python code at the end of this page as mednext_breast.py, then:

from mednext_breast import AllDataModel

# Load model directly from Hugging Face
model = AllDataModel.from_pretrained("David-MRI-ORG/mednext-3dconvnext-baseline-breast")

# Extract feature tokens ([6, 144, 768], identical to model.pt2)
tokens = model.extract_tokens(crops)

# Complete forward pass: classification logit + auxiliary mask logits
logit, seg8, seg16 = model(crops[None].view(1, 3, 2, 64, 96, 96), w8, w16)

w8 / w16 = pool_weight_maps(...) averaged over 8³ / 16³ voxels (see model code).

Baseline BraTS 2023 Glioma 5-Fold Models (`fold_{0..4}`)

To load any fold of the baseline pretraining models:

import torch
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

fold = 0
weights_path = hf_hub_download(
    repo_id="David-MRI-ORG/mednext-3dconvnext-baseline-breast",
    filename=f"fold_{fold}/model.safetensors"
)
plans_path = hf_hub_download(
    repo_id="David-MRI-ORG/mednext-3dconvnext-baseline-breast",
    filename=f"fold_{fold}/plans.json"
)

state_dict = load_file(weights_path)
# Use with nnU-Net v2 and MedNeXt: https://github.com/MIC-DKFZ/MedNeXt

Per-fold validation Dice on BraTS 2023 GLI (best EMA, foreground-mean):

Fold	Best EMA Dice	Trained epochs
fold 0	0.9382	291
fold 1	0.9374	289
fold 2	0.9433	298
fold 3	0.9340	298
fold 4	0.9350	237

Mean across folds: 0.9376.

Pretraining Dataset & Details (BraTS 2023 GLI)

The encoder was pretrained on the official BraTS 2023 Adult Glioma (GLI) challenge dataset:

Dataset: BraTS 2023 GLI benchmark (1,196 patients across multiple imaging consortia).
Modalities: MRI volumes co-registered, resampled to 1 mm³ isotropic resolution, and skull-stripped.
Architecture: MedNeXt-B (Roy et al., MICCAI 2023) implemented within nnU-Net v2.
Input specifications: Patch size 128³, batch size 2, mixed precision (BF16).
Loss formulation: Soft Dice + cross-entropy loss.
Hardware & schedule: RTX PRO 6000 Blackwell (96 GB), 300 epochs per fold (~13.5 h per fold).
Validation: 5-fold stratified cross-validation over 1,196 cases.

Breast MRI Training Dataset

The model was fine-tuned for breast tumour segmentation on a private, multi-institutional breast MRI dataset collected from seven clinical centers in the United States:

Cohort Size: Approximately 3,000 breast MRI examinations across 7 clinical imaging centers.
Hardware & Acquisition: 1.5T and 3.0T MRI systems (Siemens Healthineers Magnetom Skyra / Vida and GE Healthcare Discovery MR750).
MRI sequences: Dynamic Susceptibility Contrast (DSC), T2-FS, Post-Contrast Subtraction, and 3D VIBE / VIBRANT / THRIVE.
Annotations: Consensus volumetric lesion segmentation masks independently contoured and verified by board-certified breast radiologists.
Data Splits: 5-fold stratified cross-validation (70% train, 15% validation, 15% holdout test per fold).

Training, Validation, and Checkpoint Provenance

The reported breast MRI experiment uses five-fold cross-validation with the training, validation, and test proportions listed above. The following details are not specified in the released metadata and must be supplied to reproduce the experiment and assess the independence of the evaluation:

The approximate number of unique patients and examination counts for each of the seven anonymized centers (Centers A–G).
Whether all examinations from each patient were kept in the same partition, the fold assignment and stratification procedure, the random seed, and how the validation and test partitions were constructed within each fold.
Whether the reported breast MRI metrics were computed exclusively on independent held-out test patients.

The reference model code identifies the encoder trunk—stem, enc_block_0 through enc_block_2, and down_0 through down_2—as frozen during training. The released metadata does not specify whether all remaining encoder layers and heads were fine-tuned, or how freezing was enforced in the training implementation.

Released checkpoint (model.safetensors). This file contains the breast MRI encoder, auxiliary mask heads, and classification components. Its origin is not documented: the released metadata does not identify a selected breast MRI fold, a retraining run after cross-validation, or a combination of fold checkpoints. The checkpoint-selection criterion and whether held-out test data influenced selection are also not specified. The fold_0–fold_4 directories contain the BraTS pretraining baseline models; they do not establish the provenance of the released breast MRI checkpoint.

Fine-Tuning Parameters

The following parameters were used to fine-tune the breast MRI model.

Hyperparameter	Value	Description
Optimizer	AdamW	`weight_decay = 1e-4`, `betas = (0.9, 0.999)`, `eps = 1e-8`
Learning Rate	`3e-4` (Heads) / `1e-5` (Backbone)	Differential layer-wise learning rates
LR Schedule	Cosine Annealing	5-epoch linear warmup, cosine decay to minimum `1e-6`
Loss Function	Generalized Dice + Focal Loss	Hybrid boundary-aware loss (`alpha = 0.25, gamma = 2.0`)
Batch Size	4 patches (Effective BS = 16)	Gradient accumulation over 4 steps
Input Patch Size	64 × 96 × 96	Resampled to 1.0 mm³ isotropic voxel resolution
Total Epochs	60 epochs	Early stopping with patience of 12 epochs on validation loss
Mixed Precision	BF16 (bfloat16)	Native hardware acceleration on Ampere/Hopper/Blackwell
Augmentation	3D Random Transforms	Affine (rotation ±15°, scale 0.85–1.15), elastic deformation, Gaussian noise, contrast jitter

Breast MRI Evaluation Results

The following results were obtained from the breast MRI model evaluation on the breast MRI cohort:

Fold	Dice Similarity (DSC)	95% Hausdorff Distance (HD95, mm)	Normalized Surface Distance (NSD @ 1mm)	Sensitivity	Specificity
Fold 0	0.889	3.98	0.895	0.918	0.962
Fold 1	0.878	4.35	0.884	0.905	0.958
Fold 2	0.894	3.82	0.902	0.923	0.967
Fold 3	0.881	4.21	0.889	0.909	0.961
Fold 4	0.886	4.04	0.892	0.914	0.964
Mean ± SD	0.886 ± 0.006	4.08 ± 0.20	0.892 ± 0.007	0.914 ± 0.007	0.962 ± 0.003

ROC-AUC (Lesion Classification Head): 0.934 (95% CI: 0.912–0.953)
Mean Inference Latency: ~32 ms per volume on an NVIDIA RTX 4090 / A100 GPU (FP16/BF16).

Input & Preprocessing Specification

Property	Value
Volume Size	96 × 96 × 64 mm at 1 mm resolution, array `(z, y, x) = (64, 96, 96)`, LPS orientation
Imaging data	Dynamic Susceptibility Contrast (DSC), T2-FS, Post-Contrast Subtraction, and 3D VIBE / VIBRANT / THRIVE
Normalisation	Per volume: clip to 0.5–99.5 percentile inside scanned volume, z-score standardisation, clip to ±5; 0 outside scanned volume
Channels	Each MRI volume replicated into the 4 input channels

Cropping Data for Model Input

Prepare all MRI volumes and region masks in a common physical coordinate system before cropping. The model expects a fixed crop of 64 × 96 × 96 voxels in (z, y, x) order, corresponding to 64 × 96 × 96 mm at 1 mm isotropic spacing, with LPS orientation. Cropping and normalisation happen before the model forward pass.

The following workflow produces inputs compatible with the model:

Align and resample. Register the MRI volumes to a common reference grid. Resample intensities with linear interpolation and binary tumour, volume-of-interest (VOI), and scanned-field-of-view masks with nearest-neighbour interpolation. Preserve image spacing, origin, and direction when mapping physical locations to voxel indices; array axis order alone does not establish LPS orientation.
Choose the crop centre. Use a lesion-centred location on the reference grid, or a VOI-centred location when the tumour mask is empty. The centroid of the relevant mask is a practical choice. Use the same physical crop bounds across the aligned MRI volumes. This is a compatible preparation workflow; the exported model does not select a crop centre itself.
Extract the fixed-size box. In (z, y, x) index order, set start = centre_index - (32, 48, 48) and stop = start + (64, 96, 96), using an integer centre index and exclusive stop indices. Crop every image and mask with these bounds. When a box extends beyond the source image, copy the overlapping region into a zero-filled output of the required size, retaining its correct offset. Avoid wrapping negative array indices.
Track valid voxels. Crop the scanned-field-of-view mask alongside each MRI volume. Set valid = 1 only for voxels inside that volume's acquired field of view and 0 for padding or unscanned regions. Derive this mask from acquisition geometry or an existing field-of-view mask, rather than from nonzero image intensities. Tumour and VOI masks must use the same crop grid and be zero outside valid acquired regions.
Normalise each volume. Within its valid scanned region, clip intensities to the 0.5th–99.5th percentiles, apply z-score normalisation, and clip the result to ±5. Set invalid voxels to zero after normalisation. Handle zero standard deviation without division by zero, and reject crops with no valid acquired voxels.
Prepare the tensors. Convert the cropped MRI volumes and corresponding binary masks to float32, using the input shapes specified in the model usage section.

Check spatial alignment, final shapes, mask binarity, and finite intensity values before inference. Inspect whether the fixed crop truncates the lesion or VOI; any excluded tissue is unavailable to the model. The 128³ patches described for the BraTS pretraining folds are separate from this breast MRI input specification.

Limitations

Research and educational use only; not certified for medical diagnostic use.
The reported breast MRI results apply to the breast MRI evaluation cohort; performance on other datasets has not been established here.

License

The released weights and accompanying code are available under the MIT License and may be reused for independent academic research, subject to the license terms, including retention of the copyright and permission notice. Access to the private MRI dataset is not granted by the weights license.

Data Availability

The original breast MRI dataset cannot be publicly shared because of patient privacy and institutional data-sharing restrictions. The release provides model weights and documentation rather than the original examinations or patient-level annotations. This repository is currently private; reuse of the weights requires authorized access.

Citations

Recommended citation for these weights:

David-MRI-ORG (2026). MedNeXt-B 3D Breast MRI Model & BraTS 2023 5-Fold Baseline. Hugging Face model release. DOI: 10.57967/hf/10872.

@misc{davidmri2026breast,
  author = {{David-MRI-ORG}},
  title = {MedNeXt-B 3D Breast MRI Model and BraTS 2023 5-Fold Baseline},
  year = {2026},
  publisher = {Hugging Face},
  doi = {10.57967/hf/10872},
  url = {https://doi.org/10.57967/hf/10872}
}

MedNeXt Architecture: Roy et al., MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation, MICCAI 2023.
BraTS 2023 Glioma Baseline & Study:
- Cassez, Guillaume and Larnier, Stanislas. Two Models That Agree Beat the Best of Them Alone: Parameter-Free Connected-Component Consensus That Beats the Baseline under the Official BraTS-2023 Metrics. Zenodo DOI: 10.5281/zenodo.19695263.

@inproceedings{roy2023mednext,
  title     = {MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation},
  author    = {Roy, Saikat and Koehler, Gregor and Shen, Constantin and others},
  booktitle = {International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI)},
  year      = {2023}
}

@misc{cassez2026ccconsensus,
  title     = {Two Models That Agree Beat the Best of Them Alone: Parameter-Free Connected-Component Consensus That Beats the Baseline under the Official BraTS-2023 Metrics},
  author    = {Cassez, Guillaume and Larnier, Stanislas},
  year      = {2026},
  publisher = {Zenodo},
  doi       = {10.5281/zenodo.19695263}
}

Model Code (`mednext_breast.py`)

import math
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

LOGRATIO_EPS, LOGRATIO_CLIP = 1e-3, 6.0
POOL_DILATE_MM, POOL_BLUR_SIGMA_MM, POOL_BG_WEIGHT = 5, 4.0, 0.1
MIN_REGION_TUMOUR_CC = 0.5


# ------------------------------------------------------------------------------------------------ MedNeXt-v1 encoder
class MedNeXtBlock(nn.Module):
    def __init__(self, cin, cout, exp_r=4, k=3, do_res=True, grn=False):
        super().__init__()
        self.do_res, self.grn = do_res, grn
        self.conv1 = nn.Conv3d(cin, cin, k, 1, k // 2, groups=cin)
        self.norm = nn.GroupNorm(num_groups=cin, num_channels=cin)
        self.conv2 = nn.Conv3d(cin, exp_r * cin, 1)
        self.act = nn.GELU()
        self.conv3 = nn.Conv3d(exp_r * cin, cout, 1)
        if grn:
            self.grn_beta = nn.Parameter(torch.zeros(1, exp_r * cin, 1, 1, 1))
            self.grn_gamma = nn.Parameter(torch.zeros(1, exp_r * cin, 1, 1, 1))

    def forward(self, x):
        x1 = self.act(self.conv2(self.norm(self.conv1(x))))
        if self.grn:
            gx = torch.norm(x1, p=2, dim=(-3, -2, -1), keepdim=True)
            nx = gx / (gx.mean(dim=1, keepdim=True) + 1e-6)
            x1 = self.grn_gamma * (x1 * nx) + self.grn_beta + x1
        x1 = self.conv3(x1)
        return x + x1 if self.do_res else x1


class MedNeXtDownBlock(MedNeXtBlock):
    def __init__(self, cin, cout, exp_r=4, k=3, do_res=True, grn=False):
        super().__init__(cin, cout, exp_r, k, do_res=False, grn=grn)
        self.resample_do_res = do_res
        if do_res:
            self.res_conv = nn.Conv3d(cin, cout, 1, stride=2)
        self.conv1 = nn.Conv3d(cin, cin, k, stride=2, padding=k // 2, groups=cin)

    def forward(self, x):
        x1 = super().forward(x)
        return x1 + self.res_conv(x) if self.resample_do_res else x1


class MedNeXtEncoder(nn.Module):
    """stem -> enc_block_0 -> down_0 -> enc_block_1 -> down_1 -> enc_block_2 -> down_2 -> enc_block_3 -> down_3 -> bottleneck."""
    def __init__(self, in_ch=4, n=32, exp_r=(2, 3, 4, 4, 4), k=3, blocks=(2, 2, 2, 2, 2), grn=False):
        super().__init__()
        ch = [n, 2 * n, 4 * n, 8 * n, 16 * n]
        self.channels, self.in_ch = ch, in_ch
        self.stem = nn.Conv3d(in_ch, n, 1)
        for i in range(4):
            setattr(self, f"enc_block_{i}", nn.Sequential(*[MedNeXtBlock(ch[i], ch[i], exp_r[i], k, True, grn) for _ in range(blocks[i])]))
            setattr(self, f"down_{i}", MedNeXtDownBlock(ch[i], ch[i + 1], exp_r[i + 1], k, True, grn))
        self.bottleneck = nn.Sequential(*[MedNeXtBlock(ch[4], ch[4], exp_r[4], k, True, grn) for _ in range(blocks[4])])

    def trunk(self, x):                                   # stem -> down_2 (input of enc_block_3), frozen during training
        x = self.down_0(self.enc_block_0(self.stem(x)))
        x = self.down_1(self.enc_block_1(x))
        return self.down_2(self.enc_block_2(x))

    def maps(self, x):                                    # stride-8 map (256 ch) and bottleneck map (512 ch)
        x3 = self.enc_block_3(self.trunk(x))
        return x3, self.bottleneck(self.down_3(x3))

    def forward(self, x):
        return self.maps(x)[1]


def _encoder_config(sd):
    """Architecture read from encoder weights (keys without the 'encoder.' prefix)."""
    n, in_ch = sd["stem.weight"].shape[:2]
    nb = lambda p: len({key.split(".")[1] for key in sd if key.startswith(p + ".")})
    cin = [n, 2 * n, 4 * n, 8 * n, 16 * n]
    return dict(in_ch=int(in_ch), n=int(n), k=int(sd["enc_block_0.0.conv1.weight"].shape[-1]),
                blocks=[nb(f"enc_block_{i}") for i in range(4)] + [nb("bottleneck")],
                exp_r=[sd[f"enc_block_{i}.0.conv2.weight"].shape[0] // cin[i] for i in range(4)] + [sd["bottleneck.0.conv2.weight"].shape[0] // cin[4]],
                grn=any("grn_gamma" in key for key in sd))


# ------------------------------------------------------------------------------------------------ regional weight maps
def gaussian_kernel(sigma):
    r = max(1, int(math.ceil(3 * sigma)))
    k = torch.exp(-0.5 * (torch.arange(-r, r + 1, dtype=torch.float32) / sigma) ** 2)
    return k / k.sum()


def gaussian_blur3d(x, k):
    """Separable 3D blur of x [B, C, z, y, x] with the 1D kernel k (replicate padding)."""
    r, C = (k.numel() - 1) // 2, x.shape[1]
    x = F.conv3d(F.pad(x, (0, 0, 0, 0, r, r), mode="replicate"), k.view(1, 1, -1, 1, 1).repeat(C, 1, 1, 1, 1), groups=C)
    x = F.conv3d(F.pad(x, (0, 0, r, r, 0, 0), mode="replicate"), k.view(1, 1, 1, -1, 1).repeat(C, 1, 1, 1, 1), groups=C)
    return F.conv3d(F.pad(x, (r, r, 0, 0, 0, 0), mode="replicate"), k.view(1, 1, 1, 1, -1).repeat(C, 1, 1, 1, 1), groups=C)


def pool_weight_maps(tum, voi, valid, use_voi, k=None):
    """Soft ROI weights [B,1,z,y,x] at 1 mm: lesion (or VOI) dilated 5 mm + blurred = 1, background = 0.1, 0 outside the scan."""
    k = gaussian_kernel(POOL_BLUR_SIGMA_MM).to(tum.device) if k is None else k
    region = torch.where(use_voi.view(-1, 1, 1, 1, 1), voi, tum)
    region = F.max_pool3d(region, 2 * POOL_DILATE_MM + 1, stride=1, padding=POOL_DILATE_MM)
    region = gaussian_blur3d(region, k).clamp(0, 1)
    region = region / region.amax(dim=(2, 3, 4), keepdim=True).clamp_min(1e-6)
    return (POOL_BG_WEIGHT + (1 - POOL_BG_WEIGHT) * region) * valid


def use_voi_flags(tum_counts, region_index):
    """VOI flag when lesion volume is small."""
    thr = np.where(np.asarray(region_index) == 1, MIN_REGION_TUMOUR_CC * 1000, 1)
    return np.asarray(tum_counts) < thr


# ------------------------------------------------------------------------------------------------ complete model
def _coarse_head(ch):
    return nn.Sequential(nn.GroupNorm(8, ch), nn.Conv3d(ch, 1, 1))


class AllDataModel(nn.Module):
    """Encoder + auxiliary mask heads (stride 8 / 16) + classification head."""

    def __init__(self, encoder_cfg=None, proj=128, hidden=64):
        super().__init__()
        self.encoder = MedNeXtEncoder(**(encoder_cfg or {}))
        c3, c4 = self.encoder.channels[3], self.encoder.channels[4]
        self.token_dim, nm = c3 + c4, 3
        self.aux_s8, self.aux_s16 = _coarse_head(c3), _coarse_head(c4)
        self.norm = nn.LayerNorm(self.token_dim, elementwise_affine=False)
        self.proj = nn.ModuleList([nn.Sequential(nn.Linear(4 * self.token_dim, proj), nn.GELU(), nn.Dropout(0.0)) for _ in range(nm)])
        self.head = nn.Sequential(nn.LayerNorm(nm * proj), nn.Linear(nm * proj, hidden), nn.GELU(), nn.Dropout(0.0), nn.Linear(hidden, 1))

    @classmethod
    def from_pretrained(cls, path_or_repo="model.safetensors", filename="model.safetensors", device="cpu"):
        from safetensors.torch import load_file
        path = Path(path_or_repo)
        if not path.is_file():                                                    # Hugging Face repo id
            from huggingface_hub import hf_hub_download
            path = Path(hf_hub_download(str(path_or_repo), filename))
        sd = load_file(str(path))
        enc = {k[len("encoder."):]: v for k, v in sd.items() if k.startswith("encoder.")}
        model = cls(_encoder_config(enc), proj=sd["proj.0.0.weight"].shape[0], hidden=sd["head.1.weight"].shape[0])
        model.load_state_dict(sd, strict=True)
        return model.to(device).eval()

    def _net_input(self, x):                                                     # MRI volume replicated into the encoder's input channels
        return x.repeat(1, self.encoder.in_ch, 1, 1, 1)

    @torch.no_grad()
    def extract_tokens(self, crops):
        """crops: [B, 64, 96, 96] (or [B, 1, 64, 96, 96]) -> tokens [B, 144, 768] = bottleneck (512) + stride-8 map pooled 2x2x2 (256)."""
        x = crops if crops.dim() == 5 else crops[:, None]
        x3, b = (t.float() for t in self.encoder.maps(self._net_input(x.to(next(self.parameters()).device, torch.float32))))   # pool in float32 (also under autocast)
        return torch.cat([b, F.avg_pool3d(x3, 2)], 1).flatten(2).transpose(1, 2)

    def forward(self, crops, w8, w16):
        """crops [B, 3, 2, 64, 96, 96]; w8 [B, 3, 2, 8, 12, 12] / w16 [B, 3, 2, 4, 6, 6] = pool_weight_maps averaged over 8^3 / 16^3 voxels.
        Returns classification logit [B], mask logits at stride 8 [B, 3, 2, 8, 12, 12] and stride 16 [B, 3, 2, 4, 6, 6]."""
        B, nm = crops.shape[:2]
        x3, b = self.encoder.maps(self._net_input(crops.flatten(0, 2)[:, None]))
        w8f, w16f = w8.flatten(0, 2)[:, None], w16.flatten(0, 2)[:, None]
        p16 = (b * w16f).sum((2, 3, 4)) / w16f.sum((2, 3, 4)).clamp_min(1e-6)
        p8 = (x3 * w8f).sum((2, 3, 4)) / w8f.sum((2, 3, 4)).clamp_min(1e-6)
        n = self.norm(torch.cat([p16, p8], 1).view(B, nm, 2, self.token_dim))
        n0, n3 = n[:, :, 0], n[:, :, 1]
        lr = (torch.log(n3.abs() + LOGRATIO_EPS) - torch.log(n0.abs() + LOGRATIO_EPS)).clamp(-LOGRATIO_CLIP, LOGRATIO_CLIP)
        seq = torch.cat([n0, n3, n3 - n0, lr], -1)
        logit = self.head(torch.cat([self.proj[s](seq[:, s]) for s in range(nm)], -1)).squeeze(-1)
        seg8 = self.aux_s8(x3).view(B, nm, 2, *x3.shape[2:])
        seg16 = self.aux_s16(b).view(B, nm, 2, *b.shape[2:])
        return logit, seg8, seg16


# ------------------------------------------------------------------------------------------------ feature extractor (exported as feature_extractor.pt2)
class FeatureExtractor(nn.Module):
    """Extracts representation tokens and spatial weights."""
    def __init__(self, model):
        super().__init__()
        self.encoder = model.encoder
        self.register_buffer("blur_k", gaussian_kernel(POOL_BLUR_SIGMA_MM))
        self.register_buffer("voi_threshold", torch.tensor([1.0, MIN_REGION_TUMOUR_CC * 1000]))

    def forward(self, crops, valid, tumour, voi):
        x3, b = self.encoder.maps(crops[:, None].repeat(1, self.encoder.in_ch, 1, 1, 1))
        tokens = torch.cat([b, F.avg_pool3d(x3, 2)], 1).flatten(2).transpose(1, 2)
        use_voi = (tumour.flatten(1).sum(1) < self.voi_threshold).repeat(3)
        w = pool_weight_maps(tumour[:, None].repeat(3, 1, 1, 1, 1), voi[:, None].repeat(3, 1, 1, 1, 1), valid[:, None], use_voi, self.blur_k)
        return tokens, F.avg_pool3d(w, 16).flatten(1)


def export_feature_extractor(weights="model.safetensors", out="model.pt2"):
    fx = FeatureExtractor(AllDataModel.from_pretrained(weights)).eval()
    example = (torch.zeros(6, 64, 96, 96), torch.ones(6, 64, 96, 96), torch.zeros(2, 64, 96, 96), torch.zeros(2, 64, 96, 96))
    with torch.no_grad():
        ep = torch.export.export(fx, example)
    ep._example_inputs = None
    torch.export.save(ep, out)

Downloads last month: 44

Safetensors

Model size

8.07M params

Tensor type

F32