OSCR

3D-ViT-UNet: 3D Vision transformer based Unet-like model for Volumetric Brain Tumor Segmentation.

Code ↔ Paper

3 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 3 matches
  1. [1] § 4. Experiments › 4.4. Implementation details ↔ src/train.py, lines 128 to 169 · score 0.68 · ReduceLROnPlateau, patience, scheduler, Adam, batch, optimizer · Report a mismatch
  2. [2] § 4. Experiments › 4.2. Loss function ↔ src/train.py, lines 57 to 64 · score 0.53 · dice loss, focal loss, predicted · Report a mismatch
  3. [3] § 4. Experiments › 4.2. Loss function ↔ src/train.py, lines 57 to 64 · score 0.50 · dice loss, Focal loss · Report a mismatch

Paper

Loaded from Europe PMC by your browser; OSCR does not store it: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 173 lines · 6.1 KB · MIT · 3 matches

  1. #!/usr/bin/env python3
  2. """Train 3D ViT-UNet using preprocessed BraTS .npy volumes."""
  4. from __future__ import annotations
  6. import argparse
  7. from pathlib import Path
  9. import numpy as np
  10. import segmentation_models_pytorch as smp
  11. import torch
  12. from sklearn.metrics import accuracy_score, jaccard_score
  13. from torch import nn
  14. from torch.optim import Adam
  15. from torch.optim.lr_scheduler import ReduceLROnPlateau
  16. from torch.utils.data import DataLoader, Dataset
  17. from tqdm import tqdm
  19. try:
  20. from .model import SegNetwork
  21. except ImportError:
  22. from model import SegNetwork
  25. def dsc_from_iou(iou: float) -> float:
  26. return (2.0 * iou) / (1.0 + iou)
  29. class NpySegmentationDataset(Dataset):
  30. def __init__(self, image_dir: Path, mask_dir: Path) -> None:
  31. self.image_paths = sorted(image_dir.glob("*.npy"))
  32. self.mask_paths = sorted(mask_dir.glob("*.npy"))
  33. if not self.image_paths or not self.mask_paths:
  34. raise FileNotFoundError("No .npy files found in image/mask directories.")
  35. if len(self.image_paths) != len(self.mask_paths):
  36. raise ValueError("Image/mask count mismatch.")
  38. def __len__(self) -> int:
  39. return len(self.image_paths)
  41. def __getitem__(self, idx: int):
  42. image = np.load(self.image_paths[idx])
  43. mask = np.load(self.mask_paths[idx])
  44. if image.shape[:3] != mask.shape[:3]:
  45. raise ValueError(f"Spatial mismatch for sample index {idx}.")
  46. image_t = torch.tensor(image, dtype=torch.float32).permute(3, 0, 1, 2)
  47. mask_onehot = torch.tensor(mask, dtype=torch.float32).permute(3, 0, 1, 2)
  48. mask_indices = torch.argmax(mask_onehot, dim=0).long()
  49. return image_t, mask_indices
  52. def build_loader(image_dir: Path, mask_dir: Path, batch_size: int, shuffle: bool) -> DataLoader:
  53. dataset = NpySegmentationDataset(image_dir, mask_dir)
  54. return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=0)
  57. class CombinedLoss(nn.Module):
  58. def __init__(self) -> None:
  59. super().__init__()
  60. self.dice = smp.losses.DiceLoss(mode="multiclass")
  61. self.focal = smp.losses.FocalLoss(mode="multiclass")
  63. def forward(self, preds: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
  64. return self.dice(preds, labels) + self.focal(preds, labels)
  67. def run_epoch(
  68. model: nn.Module,
  69. loader: DataLoader,
  70. loss_fn: nn.Module,
  71. device: torch.device,
  72. optimizer: Adam | None = None,
  73. ) -> dict[str, float]:
  74. is_train = optimizer is not None
  75. model.train(is_train)
  77. total_loss, total_acc, total_iou, total_dsc = 0.0, 0.0, 0.0, 0.0
  78. loop = tqdm(loader, desc="Train" if is_train else "Val", leave=False)
  80. for images, mask_indices in loop:
  81. images = images.to(device, non_blocking=True)
  82. mask_indices = mask_indices.to(device, non_blocking=True)
  84. if is_train:
  85. optimizer.zero_grad(set_to_none=True)
  87. logits = model(images)
  88. if isinstance(logits, tuple):
  89. logits = logits[0]
  90. loss = loss_fn(logits, mask_indices)
  92. if is_train:
  93. loss.backward()
  94. optimizer.step()
  96. total_loss += loss.item()
  98. preds = torch.argmax(logits, dim=1).detach().cpu().numpy().ravel()
  99. labels = mask_indices.detach().cpu().numpy().ravel()
  100. total_acc += accuracy_score(labels, preds)
  101. iou = jaccard_score(labels, preds, average="macro")
  102. total_iou += iou
  103. total_dsc += dsc_from_iou(iou)
  105. n = max(len(loader), 1)
  106. return {
  107. "loss": total_loss / n,
  108. "acc": total_acc / n,
  109. "iou": total_iou / n,
  110. "dsc": total_dsc / n,
  111. }
  114. def parse_args() -> argparse.Namespace:
  115. parser = argparse.ArgumentParser(description="Train 3D ViT-UNet on preprocessed .npy data.")
  116. parser.add_argument("--train-images", type=Path, required=True)
  117. parser.add_argument("--train-masks", type=Path, required=True)
  118. parser.add_argument("--val-images", type=Path, required=True)
  119. parser.add_argument("--val-masks", type=Path, required=True)
  120. parser.add_argument("--epochs", type=int, default=20)
  121. parser.add_argument("--batch-size", type=int, default=1)
  122. parser.add_argument("--learning-rate", type=float, default=1e-3)
  123. parser.add_argument("--output-dir", type=Path, default=Path("checkpoints"))
  124. parser.add_argument("--resume", type=Path, default=None)
  125. return parser.parse_args()
  128. def main() -> None:
  129. args = parse_args()
  130. args.output_dir.mkdir(parents=True, exist_ok=True)
  132. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  133. print(f"Using device: {device}")
  135. train_loader = build_loader(args.train_images, args.train_masks, args.batch_size, shuffle=True)
  136. val_loader = build_loader(args.val_images, args.val_masks, args.batch_size, shuffle=False)
  138. model = SegNetwork(num_classes=4, in_chan=3).to(device)
  139. if args.resume:
  140. state = torch.load(args.resume, map_location=device)
  141. model.load_state_dict(state)
  142. print(f"Resumed from: {args.resume}")
  144. loss_fn = CombinedLoss()
  145. optimizer = Adam(model.parameters(), lr=args.learning_rate)
  146. scheduler = ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=5)
  148. best_dsc = -1.0
  149. for epoch in range(1, args.epochs + 1):
  150. train_metrics = run_epoch(model, train_loader, loss_fn, device, optimizer)
  151. val_metrics = run_epoch(model, val_loader, loss_fn, device)
  153. scheduler.step(val_metrics["loss"])
  154. lr = optimizer.param_groups[0]["lr"]
  155. print(
  156. f"Epoch {epoch}/{args.epochs} | "
  157. f"Train loss {train_metrics['loss']:.4f} DSC {train_metrics['dsc']:.4f} | "
  158. f"Val loss {val_metrics['loss']:.4f} DSC {val_metrics['dsc']:.4f} | "
  159. f"LR {lr:.6f}"
  160. )
  162. latest_path = args.output_dir / "latest.pth"
  163. torch.save(model.state_dict(), latest_path)
  165. if val_metrics["dsc"] > best_dsc:
  166. best_dsc = val_metrics["dsc"]
  167. best_path = args.output_dir / "best.pth"
  168. torch.save(model.state_dict(), best_path)
  169. print(f"Saved new best checkpoint: {best_path}")
  172. if __name__ == "__main__":
  173. main()

train.py at commit d50123d, under MIT · at the source

Overview

Authors: Sikandar Afridi1, Atif Jan1, Muhammad Abeer Irfan2, Muhammad Irfan Khattak1, Taimur Ahmed Khan1
ORCID iDs: Muhammad Abeer Irfan
  1. Department of Electrical Engineering, University of Engineering and Technology, Peshawar, Pakistan
  2. School of Computer Science, Technological University Dublin (TU Dublin), Dublin, Ireland
Institutions: University of Engineering and Technology Peshawar (Pakistan); Technological University Dublin (Ireland)
Journal: PLOS digital health, volume 5, issue 3, article e0001323
Dates: received 9 December 2025; accepted 6 March 2026; published online 31 March 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1371/journal.pdig.0001323 · PMID 41915719 · PMCID PMC13038115 · OpenAlex W7147484965
Open access: gold, a free copy (OpenAlex)
Status: code verified (how the code is found and verified)
Categories: other condition (population)
Methods: Connectivity, Machine learning
Journal subjects: Medicine and Health Sciences, Oncology, Cancers and Neoplasms, Diagnostic Medicine, Diagnostic Radiology, Magnetic Resonance Imaging, Research and Analysis Methods, Imaging Techniques, Radiology and Imaging, Malignant Tumors, Neurology, Cancer Treatment, Mathematical and Statistical Techniques, Mathematical Functions, Convolution
Topic: Advanced Neural Network Applications (Computer Vision and Pattern Recognition, Computer Science), according to OpenAlex
Citations: cited by 1 paper (Europe PMC); 66 references in the paper

Abstract

Accurate volumetric segmentation of 3D medical imaging modalities is critical for therapy planning and clinical diagnosis, particularly for brain tumor delineation. Traditional convolutional neural network (CNN)-based architectures face challenges while capturing global contextual information and modeling long-range dependencies in complex 3D volumetric data, limiting their segmentation performance. Transformer-based models have emerged as promising alternatives to CNNs for such tasks, addressing their limitations in capturing global spatial dependencies. We propose 3D-ViT-UNet, a novel U-shaped vision transformer (ViT)-based encoder-decoder architecture for end-to-end volumetric brain tumor segmentation. The model employs 3D Window Multi-Head Self-Attention (3D-W-MSA) to capture local features and a 3D Dilated-Window Multi-Head Self-Attention (3D-DW-MSA) to capture global features while reducing computational complexity. Moreover, for preserving absolute and relative positional information and preventing permutation equivalence limitation in transformers, a dynamic position encoding strategy is integrated. The proposed model demonstrates state-of-the-art (SOTA) performance for brain tumor segmentation on the BraTS 2020 dataset. It achieves a superior average Dice Similarity Coefficient (DSC) of 84.81% and a Hausdorff Distance (HD) of 4.87 mm with reduced computational complexity compared to existing methods. Also, an improvement in delineation of tumor boundaries and accurate segmentation across modalities is demonstrated through the qualitative results. Extensive quantitative and qualitative evaluations highlight the capability of 3D-ViT-UNet to achieve high accuracy with a smaller model size and lower FLOPs, making it an effective and efficient solution for clinical applications involving volumetric brain tumor segmentation.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repository

Its files are read in the Code ↔ Paper reader above, with 3 matches between paragraphs and lines of code.

maikhan/3D_ViT_Unet

License: MIT
State: the link answers, verified on 30 September 2026
Evidence: files inventoried
Commit: d50123df9a63b63bf808f1e25187f8c90e8aa9e3, 13 February 2026
Languages: Python (10)
Size: 16 files, 10 scripts
Software Heritage: not archived
Found in: “Data Availability”
Holds: README, license file, environment (requirements.txt)
Not found: CITATION.cff, tests, continuous integration, documentation
Tools: NumPy (5 files), PyTorch (3 files), scikit-learn (2 files), NiBabel (1 file)
Availability: 1 check, the latest on 30 September 2026: the link answers
  • 30 September 2026: the link answers
12 files
At the source: github.com/maikhan/3D_ViT_Unet

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 10 scripts, each with its path and the digest of its content;
  • 3 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

Datasets cited

Data Availability

The data used in this study are publicly available on Kaggle. Specifically, the volumetric brain tumor MRI data were obtained from the “Brain Tumor Segmentation (BraTS) Dataset” hosted on Kaggle and can be accessed at https://www.kaggle.com/datasets/awsaf49/brats20-dataset-training-validation All preprocessing scripts, model architecture implementations, training code, and evaluation scripts for the proposed 3D-ViT-UNet model are publicly available in a GitHub repository at https://github.com/maikhan/3D_ViT_Unet.

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 30 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 5 authors, 39 references.

Cite

This paper

Afridi, S., Jan, A., Irfan, M. A., Khattak, M. I., & Khan, T. A. (2026). 3D-ViT-UNet: 3D Vision transformer based Unet-like model for Volumetric Brain Tumor Segmentation. PLOS digital health, 5(3), e0001323. https://doi.org/10.1371/journal.pdig.0001323

BibTeX

@article{afridi20263d,
author = {Afridi, Sikandar and Jan, Atif and Irfan, Muhammad Abeer and Khattak, Muhammad Irfan and Khan, Taimur Ahmed},
title = {{3D-ViT-UNet: 3D Vision transformer based Unet-like model for Volumetric Brain Tumor Segmentation}},
journal = {PLOS digital health},
year = {2026},
month = mar,
volume = {5},
number = {3},
pages = {e0001323},
publisher = {PLOS},
issn = {2767-3170},
doi = {10.1371/journal.pdig.0001323},
url = {https://doi.org/10.1371/journal.pdig.0001323},
pmid = {41915719},
pmcid = {PMC13038115}
}

RIS

TY - JOUR
AU - Afridi, Sikandar
AU - Jan, Atif
AU - Irfan, Muhammad Abeer
AU - Khattak, Muhammad Irfan
AU - Khan, Taimur Ahmed
TI - 3D-ViT-UNet: 3D Vision transformer based Unet-like model for Volumetric Brain Tumor Segmentation
T2 - PLOS digital health
J2 - PLOS Digit Health
PY - 2026
DA - 2026/03/31
VL - 5
IS - 3
SP - e0001323
SN - 2767-3170
PB - PLOS
DO - 10.1371/journal.pdig.0001323
UR - https://doi.org/10.1371/journal.pdig.0001323
LA - en
ER -

CSL-JSON

{
"id": "10.1371/journal.pdig.0001323",
"type": "article-journal",
"title": "3D-ViT-UNet: 3D Vision transformer based Unet-like model for Volumetric Brain Tumor Segmentation",
"container-title": "PLOS digital health",
"author": [
{
"family": "Afridi",
"given": "Sikandar"
},
{
"family": "Jan",
"given": "Atif"
},
{
"family": "Irfan",
"given": "Muhammad Abeer"
},
{
"family": "Khattak",
"given": "Muhammad Irfan"
},
{
"family": "Khan",
"given": "Taimur Ahmed"
}
],
"container-title-short": "PLOS Digit Health",
"volume": "5",
"issue": "3",
"page": "e0001323",
"DOI": "10.1371/journal.pdig.0001323",
"PMID": "41915719",
"PMCID": "PMC13038115",
"ISSN": "2767-3170",
"publisher": "PLOS",
"URL": "https://doi.org/10.1371/journal.pdig.0001323",
"language": "en",
"issued": {
"date-parts": [
[
2026,
3,
31
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1038/s41598-025-09833-y [code]
A novel hybrid vision UNet architecture for brain tumor segmentation and classification.
Journal: Scientific reports
In common: NiBabel, scikit-learn, NumPy, kaggle.com/datasets/awsaf49, other condition, 3 references
[2] doi:10.2196/47445
Artificial Intelligence-Based Methods for Integrating Local and Global Features for Brain Cancer Imaging: Scoping Review.
Journal: JMIR medical informatics
In common: other condition, 7 references
[3] doi:10.21037/qims-2026-0792 [code]
An nnU-Net-based framework with adaptive feature representation for 3D brain tumor segmentation.
Journal: Quantitative imaging in medicine and surgery
In common: NiBabel, PyTorch, scikit-learn, 1 other tool, other condition, 5 references
[4] doi:10.3389/frai.2026.1807248 [code]
Federated learning framework for medical image analysis with perspective-aware contrastive and mixture of experts.
Journal: Frontiers in artificial intelligence
In common: NiBabel, PyTorch, scikit-learn, 1 other tool, kaggle.com/datasets/awsaf49, 2 references
[5] doi:10.1073/pnas.2500004122 [code]
OMT and tensor SVD-based deep learning model for segmentation and predicting genetic markers of glioma: A multicenter study.
Journal: Proceedings of the National Academy of Sciences of the United States of America
In common: NiBabel, PyTorch, scikit-learn, 1 other tool, kaggle.com/datasets/awsaf49, other condition, 2 references
[6] doi:10.3390/jimaging12060251
3D Deep Learning for Brain Tumor Segmentation and Survival Prediction: A Comprehensive Multi-Modal Analysis Using the BraTS2020 Dataset.
Journal: Journal of imaging
In common: kaggle.com/datasets/awsaf49, other condition, 4 references
[7] doi:10.3390/s25102985
DiffBTS: A Lightweight Diffusion Model for 3D Multimodal Brain Tumor Segmentation.
Journal: Sensors (Basel, Switzerland)
In common: kaggle.com/datasets/awsaf49, other condition, 4 references
[8] doi:10.3390/jimaging12070276 [code]
Hyperelastic Regularization for Near-Diffeomorphic Transformer-Based Brain MRI Registration.
Journal: Journal of imaging
In common: NiBabel, PyTorch, scikit-learn, 1 other tool, 3 references
[9] doi:10.1371/journal.pone.0304419 [code]
MCE: Medical Cognition Embedded in 3D MRI feature extraction for advancing glioma staging.
Journal: PloS one
In common: PyTorch, scikit-learn, NumPy, kaggle.com/datasets/awsaf49, other condition, 1 reference
[10] doi:10.3390/diagnostics16010139
AI-Guided Inference of Morphodynamic Attractor-like States in Glioblastoma.
Journal: Diagnostics (Basel, Switzerland)
In common: kaggle.com/datasets/awsaf49, other condition, 3 references

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Add the badge to the code's README

OSCR: the badge of this paper

If you are an author or a maintainer of this paper's code, you can add this badge. It says what OSCR knows of the paper (its code, its validated tracing map) and links the code to this page. Paste one of these snippets under the README's title. Keep the two marks around it, so that the badge can be replaced in place and is never added twice. OSCR changes nothing in your repository.

Markdown
<!-- oscr-badge:start -->
[![OSCR](https://openscicode.org/badge/doi/10.1371/journal.pdig.0001323.svg)](https://openscicode.org/badge/doi/10.1371/journal.pdig.0001323)
<!-- oscr-badge:end -->
reStructuredText
.. oscr-badge:start

.. image:: https://openscicode.org/badge/doi/10.1371/journal.pdig.0001323.svg
   :target: https://openscicode.org/badge/doi/10.1371/journal.pdig.0001323
   :alt: OSCR

.. oscr-badge:end
HTML
<!-- oscr-badge:start -->
<a href="https://openscicode.org/badge/doi/10.1371/journal.pdig.0001323"><img src="https://openscicode.org/badge/doi/10.1371/journal.pdig.0001323.svg" alt="OSCR"></a>
<!-- oscr-badge:end -->

On GitHub, open the README, choose Edit (the pencil), paste the snippet under the title, then "Commit changes".

If you choose to, you can have the badge kept up to date for you, with oscr badge --write from the command line or with the GitHub Action in your repository's workflows. Every other form and option (a card, AsciiDoc, LaTeX, shields.io) is on the badge page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: code errors and mismatches between the code and the paper about this paper and its code, from readers and its authors, beside the paper and its tracing map.

Report a mismatch between the code and the paper · Report a code error (signed in)

Reproductions: reports from readers who ran the authors' code and did not get the paper's result, what they ran, with which environment, commit and data.

Report a reproduction failure (signed in)

Versions of its code: the releases of its code tied to a version of this paper (the preprint, the accepted manuscript, the version of record…), each with the tracing map versioned with it.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.