OSCR

Distilling population specific expertise into a unified model for generalizable brain tumor segmentation.

Code ↔ Paper

17 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 17 matches · 5 of them tie a paragraph to a whole file, not to given lines: weak matches, whose lines are not tinted
  1. [1] § Materials and methods › Data preprocessing ↔ AI/data/make_dataset.py, lines 61–98 · score 0.86 · depth axis, aspect ratio, 0.65–1.5, flipping, cropped, foreground
  2. [2] § Materials and methods › Proposed MTSS-KDNet › Loss function › Logits level distillation loss ↔ AI/kd_modules/framework.py, the whole file · a weak match · score 0.85 · Binary Cross Entropy, teacher logits, distillation loss, BCE loss, Deep supervised, softens
  3. [3] § Materials and methods › Proposed MTSS-KDNet › Loss function › Supervised student loss ↔ AI/losses/loss.py, the whole file · a weak match · score 0.84 · Binary Cross Entropy, BCE loss, ground truth, Dice loss, segmentation loss, sum
  4. [4] § Materials and methods › Proposed MTSS-KDNet › Loss function › Supervised student loss ↔ AI/kd_modules/framework.py, the whole file · a weak match · score 0.84 · Binary Cross Entropy, ground truth, BCE loss, segmentation loss, deep supervision, component
  5. [5] § Materials and methods › BraTS datasets ↔ AI/losses/loss.py, the whole file · a weak match · score 0.81 · peritumoral edema, Tumor Core, Ground truth, enhancing tumor, background, TC
  6. [6] § Materials and methods › Data preprocessing ↔ AI/notebooks/Ablation/ABL_(BCE+SEG).ipynb, lines 33–180 · score 0.80 · aspect ratio, 0.65–1.5, axis, cropped, foreground, gamma
  7. [7] § Experiments and setup › Implementation details ↔ AI/notebooks/Ablation/ABL_(BCE+SEG).ipynb, lines 868–871 · score 0.79 · ReduceLROnPlateau, AdamW, weight decay, cooldown, patience, plateaus
  8. [8] § Experiments and setup › Implementation details ↔ AI/notebooks/Ablation/ABL_(KL+SEG).ipynb, lines 896–899 · score 0.79 · ReduceLROnPlateau, AdamW, weight decay, cooldown, patience, plateaus
  9. [9] § Materials and methods › Proposed MTSS-KDNet › Loss function › Latent space distillation loss ↔ AI/kd_modules/cbam_attention.py, lines 4–55 · score 0.72 · Convolutional Block Attention, feature map, emphasized, CBAM, network, Module
  10. [10] § Materials and methods › Inference ↔ AI/inference/postprocess.py, lines 5–29 · score 0.72 · sliding window inference, segmentation mask, tumor regions, prediction, model
  11. [11] § Results and discussion › Teachers progression evaluation ↔ AI/config/default_config.py, the whole file · a weak match · score 0.70 · SSA teacher, PED teacher, MEN teacher, MET teacher, teacher models, GLI
  12. [12] § Materials and methods › Inference ↔ AI/inference/predict.py, lines 37–102 · score 0.67 · sliding window inference, Gaussian, overlapping, volumes, prediction, mask
  13. [13] § Materials and methods › BraTS datasets ↔ AI/inference/preprocess.py, lines 8–49 · score 0.65 · T1 weighted, T2 weighted, FLAIR, preprocessed, volume
  14. [14] § Experiments and setup › Backbone architecture ↔ AI/models/dyn_unet.py, lines 6–20 · score 0.65 · DynUNet, medical imaging, Deep supervised, knowledge distillation, Dynamic, architecture
  15. [15] § Experiments and setup › Backbone architecture ↔ AI/models/blocks.py, lines 82–156 · score 0.63 · skip connections, medical imaging, resolutions, encoder, UNet, decoder
  16. [16] § Materials and methods › Proposed MTSS-KDNet › Teacher models ↔ AI/notebooks/Ablation/ABL_(BCE+SEG).ipynb, lines 868–871 · score 0.62 · AdamW, weight decay, Plateau, scheduler, optimized, model
  17. [17] § Results and discussion › Ablation study › Visualizing the effect of CBAM module ↔ AI/notebooks/Ablation/ABL_(KL_no_CBAM+SEG).ipynb, lines 725–784 · score 0.56 · bottleneck layer, bottleneck features, CBAM, decoder, map, ablation

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Jupyter notebook · 1,068 lines · 41 KB · no license · 3 matches

  1. # %% [markdown]
  2. # # Importing & Loading Dependencies
  4. # %%
  5. !pip install monai
  7. import nibabel as nib
  8. from monai.transforms import LoadImage, Compose, NormalizeIntensityd, RandFlipd, RandAdjustContrastd, Resized, CropForegroundd, SpatialPadd
  9. import matplotlib.pyplot as plt
  10. from torch.utils.data import DataLoader, Dataset
  11. import torch
  12. import torch.nn as nn
  13. import torch.nn.functional as F
  14. import numpy as np
  15. from typing import Optional, Sequence, Tuple, Union
  16. from torch.nn.functional import interpolate
  18. from monai.networks.blocks.convolutions import Convolution
  19. from monai.networks.layers.factories import Act, Norm
  20. from monai.networks.layers.utils import get_act_layer, get_norm_layer
  21. from monai.metrics import DiceMetric, HausdorffDistanceMetric
  23. from torch import nn, optim, amp
  24. from itertools import chain
  25. from monai.losses import DiceLoss
  27. from tqdm import tqdm
  28. from pathlib import Path
  29. import math
  30. import os
  31. import random
  33. # %% [markdown]
  34. # # Creating Dataset with Preprocessing
  36. # %%
  37. class CustomDataset3D(Dataset):
  38. def __init__(self, data_dirs, patient_lists, mode):
  39. self.data_dirs = data_dirs
  40. self.patient_lists = patient_lists
  41. self.mode = mode
  43. @staticmethod
  44. def resize_with_aspect_ratio(keys, target_size):
  45. def transform(data):
  46. for key in keys:
  47. volume = data[key]
  48. original_shape = volume.shape[-3:]
  50. scaling_factor = min(
  51. target_size[0] / original_shape[0],
  52. target_size[1] / original_shape[1],
  53. target_size[2] / original_shape[2]
  54. )
  56. # Computing the intermediate size while preserving aspect ratio
  57. new_shape = [
  58. int(dim * scaling_factor) for dim in original_shape
  59. ]
  61. # Resizing to the intermediate shape
  62. resize_transform = Resized(keys=[key], spatial_size=new_shape, mode="trilinear" if key == "imgs" else "nearest-exact")
  63. data = resize_transform(data)
  65. # Padding to the final target size
  66. pad_transform = SpatialPadd(keys=[key], spatial_size=target_size, mode="constant")
  67. data = pad_transform(data)
  68. return data
  70. return transform
  72. def preprocess(cls, data, mode):
  73. if mode == 'training':
  74. transform = Compose([
  75. CropForegroundd(keys=["imgs", "masks"], source_key="imgs"),
  76. cls.resize_with_aspect_ratio(keys=["imgs", "masks"], target_size=[128, 128, 128]),
  77. NormalizeIntensityd( keys=['imgs'], nonzero=False, channel_wise=True),
  79. RandFlipd(keys=["imgs", "masks"],
  80. prob=0.5,
  81. spatial_axis=2,
  82. ),
  84. RandAdjustContrastd(
  85. keys=["imgs"],
  86. prob=0.15,
  87. gamma=(0.65, 1.5),
  88. ),
  90. ])
  92. elif mode == 'validation':
  93. transform = Compose([
  94. CropForegroundd(keys=["imgs", "masks"], source_key="imgs"),
  95. cls.resize_with_aspect_ratio(keys=["imgs", "masks"], target_size=[128, 128, 128]),
  96. NormalizeIntensityd( keys=['imgs'], nonzero=False, channel_wise=True)
  98. ])
  100. else: # 'testing'
  101. transform = Compose([
  102. CropForegroundd(keys=["imgs", "masks"], source_key="imgs"),
  103. cls.resize_with_aspect_ratio(keys=["imgs", "masks"], target_size=[128, 128, 128]),
  104. NormalizeIntensityd( keys=['imgs'], nonzero=False, channel_wise=True)
  106. ])
  108. augmented_data = transform(data)
  109. return augmented_data
  111. def __len__(self):
  112. return len(self.patient_lists)
  114. def __getitem__(self, idx):
  115. patient_id = self.patient_lists[idx]
  116. loadimage = LoadImage(reader='NibabelReader', image_only=True)
  118. data_type=patient_id.split('-')[1]
  119. if data_type == 'GLI':
  120. patient_folder_path = os.path.join('/kaggle/input/bratsglioma/Training', patient_id)
  121. elif data_type == 'SSA':
  122. patient_folder_path = os.path.join('/kaggle/input/bratsafrica24', patient_id)
  123. elif data_type == 'PED':
  124. patient_folder_path = os.path.join('/kaggle/input/bratsped/Training', patient_id)
  125. elif data_type == 'MEN':
  126. patient_folder_path = os.path.join('/kaggle/input/bratsmen', patient_id)
  127. else:
  128. patient_folder_path = os.path.join('/kaggle/input/bratsmet24', patient_id)
  130. def resolve_file_path(folder, name):
  131. file_path = os.path.join(folder, name)
  132. # Check if the given path is a directory (case with 4 subdirs)
  133. if os.path.isdir(file_path):
  134. # Find the first file inside the directory that ends with .nii
  135. for root, _, files in os.walk(file_path):
  136. for file in files:
  137. if file.endswith(".nii"):
  138. return os.path.join(root, file)
  139. return file_path
  141. # Resolve paths for all required image types
  142. t1c_path = resolve_file_path(patient_folder_path, patient_id + '-t1c.nii')
  143. t1n_path = resolve_file_path(patient_folder_path, patient_id + '-t1n.nii')
  144. t2f_path = resolve_file_path(patient_folder_path, patient_id + '-t2f.nii')
  145. t2w_path = resolve_file_path(patient_folder_path, patient_id + '-t2w.nii')
  146. seg_path = os.path.join(patient_folder_path, patient_id + '-seg.nii')
  148. t1c_loader = loadimage( t1c_path )
  149. t1n_loader = loadimage( t1n_path )
  150. t2f_loader = loadimage( t2f_path )
  151. t2w_loader = loadimage( t2w_path )
  152. masks_loader = loadimage( seg_path )
  154. # Make the dimension of channel
  155. t1c_tensor = torch.Tensor(t1c_loader).unsqueeze(0)
  156. t1n_tensor = torch.Tensor(t1n_loader).unsqueeze(0)
  157. t2f_tensor = torch.Tensor(t2f_loader).unsqueeze(0)
  158. t2w_tensor = torch.Tensor(t2w_loader).unsqueeze(0)
  159. masks_tensor = torch.Tensor(masks_loader).unsqueeze(0)
  161. concat_tensor = torch.cat( (t1c_tensor, t1n_tensor, t2f_tensor, t2w_tensor, masks_tensor), 0 )
  162. data = {
  163. 'imgs' : np.array(concat_tensor[0:4,:,:,:]),
  164. 'masks' : np.array(concat_tensor[4:,:,:,:])
  165. }
  167. augmented_imgs_masks = self.preprocess(data, self.mode)
  168. imgs = np.array(augmented_imgs_masks['imgs'])
  169. masks = np.array(augmented_imgs_masks['masks'])
  171. y = {
  173. 'imgs' : torch.from_numpy(imgs).type(torch.FloatTensor),
  174. 'masks' : torch.from_numpy(masks).type(torch.FloatTensor),
  175. 'patient_id' : patient_id,
  176. 'data_type' : data_type
  178. }
  180. return y
  182. # %% [markdown]
  183. # # Data Loaders
  185. # %%
  186. def combine_datasets(dataset_lists, batch_size=3):
  187. max_len = max(len(dataset) for dataset in dataset_lists)
  189. # Ensure batch_size matches the number of datasets
  190. if batch_size != len(dataset_lists):
  191. raise ValueError("Batch size must equal the number of datasets for this function.")
  193. combined_paths = []
  195. for i in range(0, max_len, batch_size):
  196. for j in range(batch_size):
  197. index = (i + j) % max_len
  198. batch = [dataset[index % len(dataset)] for dataset in dataset_lists]
  199. combined_paths.extend(batch)
  201. # if j == 0:
  202. # print(f"Batch {(i // batch_size) + 1}: {batch}")
  204. return combined_paths
  206. # %%
  207. def prepare_data_loaders(args):
  208. train_datasets, val_datasets, test_datasets = [], [], []
  209. split_ratio = {'training': 0.71, 'validation': 0.09, 'testing': 0.2}
  211. for i, data_dir in enumerate(args['data_dirs']):
  212. patient_lists = os.listdir( data_dir )
  213. patient_lists.sort()
  214. total_patients = len(patient_lists)
  216. random.seed(5)
  217. random.shuffle(patient_lists)
  219. train_split = int(split_ratio['training'] * total_patients)
  220. val_split = int(split_ratio['validation'] * total_patients)
  222. train_patient_lists = patient_lists[:train_split]
  223. val_patient_lists = patient_lists[train_split : train_split + val_split]
  224. test_patient_lists = patient_lists[train_split + val_split :]
  226. train_patient_lists.sort()
  227. val_patient_lists.sort()
  228. test_patient_lists.sort()
  230. print(f'Number of training samples in {data_dir.split("/")[3]} DataSet: {len(train_patient_lists)}')
  231. print(f'Number of validation samples in {data_dir.split("/")[3]} DataSet: {len(val_patient_lists)}')
  232. print(f'Number of testing samples in {data_dir.split("/")[3]} DataSet: {len(test_patient_lists)} ')
  234. train_datasets.append(train_patient_lists)
  235. val_datasets.append(val_patient_lists)
  236. test_datasets.append(test_patient_lists)
  238. combined_trainDataset = combine_datasets(train_datasets, batch_size=args['train_batch_size'])
  239. combined_valDataset = list(chain.from_iterable(val_datasets))
  240. combined_testDataset = list(chain.from_iterable(test_datasets))
  242. print(f'Number of combined training samples', len(combined_trainDataset))
  243. print(f'Number of combined validation samples', len(combined_valDataset))
  244. print(f'Number of combined testing samples', len(combined_testDataset))
  246. trainDataset = CustomDataset3D( args['data_dirs'], combined_trainDataset, mode='training')
  247. valDataset = CustomDataset3D( args['data_dirs'], combined_valDataset, mode='validation')
  248. testDataset = CustomDataset3D( args['data_dirs'], combined_testDataset, mode='testing')
  250. trainLoader = DataLoader(
  251. trainDataset, batch_size=args['train_batch_size'], num_workers=args['workers'], prefetch_factor=2,
  252. pin_memory=True, shuffle=False)
  254. valLoader = DataLoader(
  255. valDataset, batch_size=args['val_batch_size'], num_workers=args['workers'], prefetch_factor=2,
  256. pin_memory=True, shuffle=False)
  258. testLoader = DataLoader(
  259. testDataset, batch_size=args['test_batch_size'], num_workers=args['workers'], prefetch_factor=2,
  260. pin_memory=True, shuffle=False)
  262. return trainLoader, valLoader, testLoader
  264. # %% [markdown]
  265. # # Visualizing Data
  267. # %%
  268. # args = {
  269. # 'workers': 2,
  270. # 'epochs': 10,
  271. # 'train_batch_size': 2,
  272. # 'val_batch_size': 2,
  273. # 'test_batch_size': 2,
  274. # 'learning_rate': 1e-3,
  275. # 'weight_decay': 1e-5,
  276. # 'lambd': 0.0051,
  277. # 'data_dir': '/kaggle/input/bratsafrica24/',
  278. # 'in_checkpoint_dir': Path('/kaggle/input/adultgliomamodel-45epochs'),
  279. # 'out_checkpoint_dir': Path('/kaggle/working/')
  280. # }
  282. # trainLoader, valLoader, testLoader = prepare_data_loaders(args)
  284. # for step, y in enumerate( trainLoader ):
  285. # print(y['imgs'].shape)
  286. # print(y['patient_id'])
  288. # fig, axes = plt.subplots(1, 4, figsize=(16, 4))
  289. # for sequence in range(4):
  290. # sequence_data = y['imgs'][0][sequence, :, :, :].cpu().detach().numpy()
  291. # slice_index = sequence_data.shape[2] // 2
  292. # axes[sequence].imshow(np.rot90(sequence_data[:, :, slice_index]), cmap='gray', origin='lower')
  293. # axes[sequence].set_title(f'Sequence {sequence + 1}')
  295. # plt.show()
  297. # %% [markdown]
  298. # # DynUNet Model
  300. # %%
  301. class UnetBasicBlock(nn.Module):
  302. """
  303. A CNN module module that can be used for DynUNet, based on:
  304. `Automated Design of Deep Learning Methods for Biomedical Image Segmentation <https://arxiv.org/abs/1904.08128>`_.
  305. `nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation <https://arxiv.org/abs/1809.10486>`_.
  307. Args:
  308. spatial_dims: number of spatial dimensions.
  309. in_channels: number of input channels.
  310. out_channels: number of output channels.
  311. kernel_size: convolution kernel size.
  312. stride: convolution stride.
  313. norm_name: feature normalization type and arguments.
  314. act_name: activation layer type and arguments.
  315. dropout: dropout probability.
  317. """
  319. def __init__(
  320. self,
  321. spatial_dims: int,
  322. in_channels: int,
  323. out_channels: int,
  324. kernel_size: Union[Sequence[int], int],
  325. stride: Union[Sequence[int], int],
  326. norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}),
  327. act_name: Union[Tuple, str] = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
  328. dropout: Optional[Union[Tuple, str, float]] = None,
  329. ):
  330. super().__init__()
  331. self.conv1 = get_conv_layer(
  332. spatial_dims,
  333. in_channels,
  334. out_channels,
  335. kernel_size=kernel_size,
  336. stride=stride,
  337. dropout=dropout,
  338. conv_only=True,
  339. )
  341. self.conv2 = get_conv_layer(
  342. spatial_dims,
  343. out_channels,
  344. out_channels,
  345. kernel_size=kernel_size,
  346. stride=1,
  347. dropout=dropout,
  348. conv_only=True
  349. )
  350. self.lrelu = get_act_layer(name=act_name)
  351. self.norm1 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels)
  352. self.norm2 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels)
  354. def forward(self, inp):
  355. out = self.conv1(inp)
  356. out = self.norm1(out)
  357. out = self.lrelu(out)
  358. out = self.conv2(out)
  359. out = self.norm2(out)
  360. out = self.lrelu(out)
  361. return out
  365. class UnetUpBlock(nn.Module):
  366. """
  367. An upsampling module that can be used for DynUNet, based on:
  368. `Automated Design of Deep Learning Methods for Biomedical Image Segmentation <https://arxiv.org/abs/1904.08128>`_.
  369. `nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation <https://arxiv.org/abs/1809.10486>`_.
  371. Args:
  372. spatial_dims: number of spatial dimensions.
  373. in_channels: number of input channels.
  374. out_channels: number of output channels.
  375. kernel_size: convolution kernel size.
  376. stride: convolution stride.
  377. upsample_kernel_size: convolution kernel size for transposed convolution layers.
  378. norm_name: feature normalization type and arguments.
  379. act_name: activation layer type and arguments.
  380. dropout: dropout probability.
  381. trans_bias: transposed convolution bias.
  383. """
  385. def __init__(
  386. self,
  387. spatial_dims: int,
  388. in_channels: int,
  389. out_channels: int,
  390. kernel_size: Union[Sequence[int], int],
  391. upsample_kernel_size: Union[Sequence[int], int],
  392. norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}),
  393. act_name: Union[Tuple, str] = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
  394. dropout: Optional[Union[Tuple, str, float]] = None,
  395. trans_bias: bool = False,
  396. ):
  397. super().__init__()
  398. upsample_stride = upsample_kernel_size
  400. # ( a purple arrow in the paper )
  401. self.transp_conv = get_conv_layer(
  402. spatial_dims,
  403. in_channels,
  404. out_channels,
  405. kernel_size=upsample_kernel_size,
  406. stride=upsample_stride,
  407. dropout=dropout,
  408. bias=trans_bias,
  409. conv_only=True,
  410. is_transposed=True,
  411. )
  413. # A light blue conv blocks in the decoder of nnUNet
  414. self.conv_block = UnetBasicBlock(
  415. spatial_dims,
  416. out_channels + out_channels,
  417. out_channels,
  418. kernel_size=kernel_size,
  419. stride=1,
  420. dropout=dropout,
  421. norm_name=norm_name,
  422. act_name=act_name,
  423. )
  425. def forward(self, inp, skip):
  426. # number of channels for skip should equals to out_channels
  427. out = self.transp_conv(inp)
  428. out = torch.cat((out, skip), dim=1)
  429. out = self.conv_block(out)
  430. return out
  434. class UnetOutBlock(nn.Module):
  435. def __init__(
  436. self, spatial_dims: int, in_channels: int, out_channels: int, dropout: Optional[Union[Tuple, str, float]] = None
  437. ):
  438. super().__init__()
  439. self.conv = get_conv_layer(
  440. spatial_dims, in_channels, out_channels, kernel_size=1, stride=1, dropout=dropout, bias=True, conv_only=True
  441. )
  443. def forward(self, inp):
  444. return self.conv(inp)
  447. def get_conv_layer(
  448. spatial_dims: int,
  449. in_channels: int,
  450. out_channels: int,
  451. kernel_size: Union[Sequence[int], int] = 3,
  452. stride: Union[Sequence[int], int] = 1,
  453. act: Optional[Union[Tuple, str]] = Act.PRELU,
  454. norm: Union[Tuple, str] = Norm.INSTANCE,
  455. dropout: Optional[Union[Tuple, str, float]] = None,
  456. bias: bool = False,
  457. conv_only: bool = True,
  458. is_transposed: bool = False,
  459. ):
  460. padding = get_padding(kernel_size, stride)
  461. output_padding = None
  462. if is_transposed:
  463. output_padding = get_output_padding(kernel_size, stride, padding)
  465. return Convolution(
  466. spatial_dims,
  467. in_channels,
  468. out_channels,
  469. strides=stride,
  470. kernel_size=kernel_size,
  471. act=act,
  472. norm=norm,
  473. dropout=dropout,
  474. bias=bias,
  475. conv_only=conv_only,
  476. is_transposed=is_transposed,
  477. padding=padding,
  478. output_padding=output_padding,
  479. )
  482. def get_padding(
  483. kernel_size: Union[Sequence[int], int], stride: Union[Sequence[int], int]
  484. ) -> Union[Tuple[int, ...], int]:
  486. kernel_size_np = np.atleast_1d(kernel_size)
  487. stride_np = np.atleast_1d(stride)
  488. padding_np = (kernel_size_np - stride_np + 1) / 2
  489. if np.min(padding_np) < 0:
  490. raise AssertionError("padding value should not be negative, please change the kernel size and/or stride.")
  491. padding = tuple(int(p) for p in padding_np)
  493. return padding if len(padding) > 1 else padding[0]
  496. def get_output_padding(
  497. kernel_size: Union[Sequence[int], int], stride: Union[Sequence[int], int], padding: Union[Sequence[int], int]
  498. ) -> Union[Tuple[int, ...], int]:
  499. kernel_size_np = np.atleast_1d(kernel_size)
  500. stride_np = np.atleast_1d(stride)
  501. padding_np = np.atleast_1d(padding)
  503. out_padding_np = 2 * padding_np + stride_np - kernel_size_np
  504. if np.min(out_padding_np) < 0:
  505. raise AssertionError("out_padding value should not be negative, please change the kernel size and/or stride.")
  506. out_padding = tuple(int(p) for p in out_padding_np)
  508. return out_padding if len(out_padding) > 1 else out_padding[0]
  510. def set_requires_grad(nets, requires_grad=False):
  511. if not isinstance(nets, list):
  512. nets = [nets]
  513. for net in nets:
  514. if net is not None:
  515. for param in net.parameters():
  516. param.requires_grad = requires_grad
  518. # %%
  519. class DynUNet(nn.Module):
  520. def __init__(
  521. self,
  522. spatial_dims: int,
  523. in_channels: int,
  524. out_channels: int,
  525. deep_supervision: bool,
  526. KD: bool = False
  527. ):
  528. super().__init__()
  529. self.spatial_dims = spatial_dims
  530. self.in_channels = in_channels
  531. self.out_channels = out_channels
  532. self.deep_supervision = deep_supervision
  533. self.KD_enabled = KD
  535. self.input_conv = UnetBasicBlock( spatial_dims=self.spatial_dims,
  536. in_channels=self.in_channels,
  537. out_channels=64,
  538. kernel_size=3,
  539. stride=1
  540. )
  541. self.down1 = UnetBasicBlock( spatial_dims=self.spatial_dims,
  542. in_channels=64,
  543. out_channels=96,
  544. kernel_size=3,
  545. stride=2 # Reduces spatial dims by 2
  546. )
  547. self.down2 = UnetBasicBlock( spatial_dims=self.spatial_dims,
  548. in_channels=96,
  549. out_channels=128,
  550. kernel_size=3,
  551. stride=2
  552. )
  553. self.down3 = UnetBasicBlock( spatial_dims=self.spatial_dims,
  554. in_channels=128,
  555. out_channels=192,
  556. kernel_size=3,
  557. stride=2
  558. )
  559. self.down4 = UnetBasicBlock( spatial_dims=self.spatial_dims,
  560. in_channels=192,
  561. out_channels=256,
  562. kernel_size=3,
  563. stride=2
  564. )
  565. self.down5 = UnetBasicBlock( spatial_dims=self.spatial_dims,
  566. in_channels=256,
  567. out_channels=384,
  568. kernel_size=3,
  569. stride=2
  570. )
  571. self.bottleneck = UnetBasicBlock( spatial_dims=self.spatial_dims,
  572. in_channels=384,
  573. out_channels=512,
  574. kernel_size=3,
  575. stride=2
  576. )
  577. self.up1 = UnetUpBlock( spatial_dims=self.spatial_dims,
  578. in_channels=512,
  579. out_channels=384,
  580. kernel_size=3,
  581. upsample_kernel_size=2
  582. )
  583. self.up2 = UnetUpBlock( spatial_dims=self.spatial_dims,
  584. in_channels=384,
  585. out_channels=256,
  586. kernel_size=3,
  587. upsample_kernel_size=2
  588. )
  589. self.up3 = UnetUpBlock( spatial_dims=self.spatial_dims,
  590. in_channels=256,
  591. out_channels=192,
  592. kernel_size=3,
  593. upsample_kernel_size=2
  594. )
  595. self.up4 = UnetUpBlock( spatial_dims=self.spatial_dims,
  596. in_channels=192,
  597. out_channels=128,
  598. kernel_size=3,
  599. upsample_kernel_size=2
  600. )
  602. self.up5 = UnetUpBlock( spatial_dims=self.spatial_dims,
  603. in_channels=128,
  604. out_channels=96,
  605. kernel_size=3,
  606. upsample_kernel_size=2
  607. )
  608. self.up6 = UnetUpBlock( spatial_dims=self.spatial_dims,
  609. in_channels=96,
  610. out_channels=64,
  611. kernel_size=3,
  612. upsample_kernel_size=2
  613. )
  614. self.out1 = UnetOutBlock( spatial_dims=self.spatial_dims,
  615. in_channels=64,
  616. out_channels=self.out_channels,
  617. )
  618. self.out2 = UnetOutBlock( spatial_dims=self.spatial_dims,
  619. in_channels=96,
  620. out_channels=self.out_channels,
  621. )
  622. self.out3 = UnetOutBlock( spatial_dims=self.spatial_dims,
  623. in_channels=128,
  624. out_channels=self.out_channels,
  625. )
  627. def forward( self, input ):
  629. # Input
  630. x0 = self.input_conv( input ) # x0.shape = (B x 64 x 128 x 128 x 128)
  632. # Encoder
  633. x1 = self.down1( x0 ) # x1.shape = (B x 96 x 64 x 64 x 64)
  634. x2 = self.down2( x1 ) # x2.shape = (B x 128 x 32 x 32 x 32)
  635. x3 = self.down3( x2 ) # x3.shape = (B x 192 x 16 x 16 x 16)
  636. x4 = self.down4( x3 ) # x4.shape = (B x 256 x 8 x 8 x 8)
  637. x5 = self.down5( x4 ) # x5.shape = (B x 384 x 4 x 4 x 4)
  639. # Bottleneck
  640. x6 = self.bottleneck( x5 ) # x6.shape = (B x 512 x 2 x 2 x 2)
  642. # Decoder
  643. x7 = self.up1( x6, x5 ) # x7.shape = (B x 384 x 4 x 4 x 4)
  644. x8 = self.up2( x7, x4 ) # x8.shape = (B x 256 x 8 x 8 x 8)
  645. x9 = self.up3( x8, x3 ) # x9.shape = (B x 192 x 16 x 16 x 16)
  646. x10 = self.up4( x9, x2 ) # x10.shape = (B x 128 x 32 x 32 x 32)
  647. x11 = self.up5( x10, x1 ) # x11.shape = (B x 96 x 64 x 64 x 64)
  648. x12 = self.up6( x11, x0 ) # x12.shape = (B x 64 x 128 x 128 x 128)
  650. # Output
  651. output1 = self.out1( x12 )
  653. if (self.training and self.deep_supervision) or self.KD_enabled:
  655. # output['pred'].shape = B x 3 x 4 x 128 x 128 x 128
  656. output2 = interpolate( self.out2( x11 ), output1.shape[2:])
  657. output3 = interpolate( self.out3( x10 ), output1.shape[2:])
  658. output_all = [ output1, output2, output3 ]
  659. return { 'pred' : torch.stack(output_all, dim=1),
  660. 'bottleneck_feature_map' : x6 }
  662. return { 'pred' : output1 }
  664. # %% [markdown]
  665. # # Visualizing Model Instance
  667. # %%
  668. # !pip install torchsummary
  669. # from torchsummary import summary
  671. # # Initialize your DynUNet model
  672. # model = DynUNet(spatial_dims=3, in_channels=4, out_channels=4, deep_supervision=True, KD=True)
  674. # # Print model summary
  675. # summary(model, input_size=(4, 128, 128, 128)) # Adjust input_size according to your needs
  677. # %% [markdown]
  678. # # ClearML
  680. # %%
  681. !pip install clearml
  682. from clearml import Task
  684. %env CLEARML_WEB_HOST=https://app.clear.ml/
  685. %env CLEARML_API_HOST=https://api.clear.ml
  686. %env CLEARML_FILES_HOST=https://files.clear.ml
  687. %env CLEARML_API_ACCESS_KEY=CLEARML_API_ACCESS_KEY
  688. %env CLEARML_API_SECRET_KEY=CLEARML_API_SECRET_KEY
  690. # %% [markdown]
  691. # # GPUs Check
  693. # %%
  694. if torch.cuda.is_available():
  695. num_gpus = torch.cuda.device_count()
  696. print(f"Number of GPUs available: {num_gpus}")
  697. for i in range(num_gpus):
  698. print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
  699. else:
  700. print("No GPU available. Running on CPU.")
  702. # %%
  703. # # For freeing gpu
  704. # import gc; gc.collect(); torch.cuda.empty_cache()
  706. # %% [markdown]
  707. # # Loss Function
  709. # %%
  710. class LossFunction(nn.Module):
  711. def __init__(self):
  712. super(LossFunction, self).__init__()
  713. self.dice = DiceLoss(sigmoid=True, batch=True, smooth_nr=1e-05, smooth_dr=1e-05)
  714. self.ce = nn.BCEWithLogitsLoss()
  716. def _loss(self, p, y):
  717. return self.dice(p, y) + self.ce(p, y.float())
  719. def forward(self, p, y):
  720. y_wt, y_tc, y_et = y > 0, ((y == 1) + (y == 3)) > 0, y == 3
  721. p_wt, p_tc, p_et = p[:, 1].unsqueeze(1), p[:, 2].unsqueeze(1), p[:, 3].unsqueeze(1)
  722. l_wt, l_tc, l_et = self._loss(p_wt, y_wt), self._loss(p_tc, y_tc), self._loss(p_et, y_et)
  723. return l_wt + l_tc + l_et
  725. # %% [markdown]
  726. # # Student KD Model
  728. # %%
  729. class Student_KD_loss(nn.Module):
  730. def __init__(self):
  731. super().__init__()
  732. self.student = DynUNet( spatial_dims=3, in_channels=4, out_channels=4, deep_supervision=True)
  733. self.loss_fn = LossFunction()
  734. self.temperature = 5.0
  735. self.bce_loss = nn.BCEWithLogitsLoss()
  737. def forward(self, teacher_outputs, y):
  738. with amp.autocast('cuda:1'):
  739. student_outputs = self.student( y['imgs'] )
  741. # Student loss with Deep supervision -> (Dice loss)
  742. segloss_s_decoder_1 = self.loss_fn( student_outputs['pred'][:,0], y['masks'] ) # student_outputs['pred'].shape = B x 3 x 4 x 128 x 128 x 128
  743. segloss_s_decoder_2 = self.loss_fn( student_outputs['pred'][:,1], y['masks'] )
  744. segloss_s_decoder_3 = self.loss_fn( student_outputs['pred'][:,2], y['masks'] )
  746. student_seg_loss = segloss_s_decoder_1 + 0.5*segloss_s_decoder_2 + 0.25*segloss_s_decoder_3
  748. # KD loss between prediction layers -> (BCE Loss on logits with DS)
  749. bce_loss_with_teacher = 0
  750. decoder_weights = [1, 0.5, 0.25]
  752. for decoder_idx, weight in enumerate(decoder_weights):
  753. teacher_logits = teacher_outputs['pred'][:, decoder_idx] / self.temperature # Shape: (B, 4, 128, 128, 128)
  754. student_logits = student_outputs['pred'][:, decoder_idx] / self.temperature # Shape: (B, 4, 128, 128, 128)
  756. # Compute hard labels for WT, TC, ET
  757. teacher_probs = torch.sigmoid(teacher_logits)
  758. teacher_hard_labels = [
  759. (teacher_probs[:, channel_idx] > 0.5).float().unsqueeze(1) for channel_idx in range(1, 4)
  760. ]
  762. # Extract student logits for WT, TC, ET
  763. student_logits_channels = [
  764. student_logits[:, channel_idx].unsqueeze(1) for channel_idx in range(1, 4)
  765. ]
  767. # Compute BCE losses for WT, TC, ET
  768. bce_losses = [
  769. self.bce_loss(student_logits_channels[i], teacher_hard_labels[i]) for i in range(3)
  770. ]
  772. # Aggregate loss for this decoder with the corresponding weight
  773. bce_loss_with_teacher += sum(bce_losses) * (self.temperature ** 2) * weight
  775. #-----------------------------------------------------------------------------------#
  777. alpha, zeta = 1.0, 0.1
  779. print("Seg loss: ", student_seg_loss)
  780. print("BCE loss with teacher: ", bce_loss_with_teacher)
  782. print("Seg loss weighted: ", alpha*student_seg_loss)
  783. print("BCE loss with teacher weighted: ", zeta*bce_loss_with_teacher)
  785. batch_total_student_loss = alpha*student_seg_loss + zeta*bce_loss_with_teacher
  787. print("-------------Final student loss-------------")
  788. print(batch_total_student_loss)
  789. print("-------------Final student loss-------------")
  791. KD_output = {
  792. 'batch_total_student_loss' : batch_total_student_loss,
  793. 'seg_weighted' : alpha*student_seg_loss,
  794. 'bce_weighted' : zeta*bce_loss_with_teacher,
  795. }
  797. return KD_output
  799. # %% [markdown]
  800. # # Training & Validation
  802. # %%
  803. def evaluate(model, loader, epoch, task):
  804. torch.manual_seed(0)
  805. model.eval()
  806. loss_fn = LossFunction()
  807. n_val_batches = len(loader)
  809. tumors_val_losses, running_loss = validate_model(model, loader, loss_fn)
  810. epoch_val_loss = running_loss / n_val_batches
  811. log_val_epoch_losses(tumors_val_losses, epoch, task, epoch_val_loss)
  813. print(f"------Final validation dice loss after epoch {epoch + 1}: {epoch_val_loss}-------")
  815. model.student.to('cuda:1')
  816. model.train()
  818. return epoch_val_loss
  820. def validate_model(model, loader, loss_fn):
  821. tumors_val_losses = {'GLI': [], 'PED': [], 'SSA': [], 'MEN':[], 'MET':[]}
  822. running_loss = 0
  823. n_val_batches = len(loader)
  825. with tqdm(total=n_val_batches, desc='Validating', unit='batch', leave=False) as pbar:
  826. with torch.no_grad():
  827. for y in loader:
  828. val_loss, data_type = process_batch(model, y, loss_fn)
  829. tumors_val_losses[data_type].append(val_loss.item())
  830. running_loss += val_loss
  831. pbar.update(1)
  833. return tumors_val_losses, running_loss
  835. def process_batch(model, y, loss_fn):
  836. y['imgs'], y['masks'] = y['imgs'].to('cuda'), y['masks'].to('cuda')
  837. data_type = y['data_type'][0]
  839. with torch.amp.autocast('cuda'):
  840. output = model.student.to('cuda')(y['imgs'])
  841. val_loss = loss_fn(output['pred'], y['masks'])
  843. print(f"Validation dice loss per batch: {val_loss}")
  844. return val_loss, data_type
  846. def log_val_epoch_losses(tumors_val_losses, epoch, task, epoch_val_loss):
  847. for tumor_type, losses in tumors_val_losses.items():
  848. avg_loss = sum(losses) / len(losses) if losses else 0
  849. task.get_logger().report_scalar(
  850. title=f"{tumor_type} Losses over Epochs",
  851. series=f"{tumor_type} Epoch valLoss",
  852. iteration=epoch + 1,
  853. value=avg_loss
  854. )
  856. task.get_logger().report_scalar("KD Losses over Epochs", "val_loss", iteration=epoch+1, value=epoch_val_loss)
  858. # %%
  859. def setup_environment(args):
  860. torch.manual_seed(0)
  861. args['out_checkpoint_dir'].mkdir(parents=True, exist_ok=True)
  863. def initialize_models():
  864. teacher_model = DynUNet(spatial_dims=3, in_channels=4, out_channels=4, deep_supervision=True, KD=True).to('cuda:0')
  865. student_model = Student_KD_loss().to('cuda:1')
  866. return teacher_model, student_model
  868. def initialize_optimizer_scheduler(student_model, args):
  869. optimizer = optim.AdamW(student_model.parameters(), lr=args['learning_rate'], weight_decay=args['weight_decay'], eps=1e-4)
  870. scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, cooldown=1, threshold=0.001, min_lr=1e-6)
  871. return optimizer, scheduler
  873. def load_teacher_model(teacher_model, data_type, teacher_model_paths):
  874. teacher_model_path = teacher_model_paths.get(data_type)
  875. if teacher_model_path and Path(teacher_model_path).is_file():
  876. ckpt = torch.load(teacher_model_path, map_location='cuda:0', weights_only=True)
  877. teacher_model.load_state_dict(ckpt['teacher_model'])
  878. print(f"Loaded model: {teacher_model_path}")
  880. def load_student_checkpoint(student_model, optimizer, scaler, scheduler, args):
  881. checkpoint_path = args['in_checkpoint_dir'] / 'Student_model_after_epoch_58_trainLoss_0.8541_valLoss_0.3335.pth'
  882. if checkpoint_path.is_file():
  883. print(f"Found model {checkpoint_path}")
  884. ckpt = torch.load(checkpoint_path, map_location='cuda:1', weights_only=True)
  885. student_model.student.load_state_dict(ckpt['student_model'])
  886. optimizer.load_state_dict(ckpt['optimizer_student'])
  887. scaler.load_state_dict(ckpt['grad_scaler_state'])
  888. scheduler.load_state_dict(ckpt['scheduler_state_dict'])
  889. print(f"Loaded student model: {checkpoint_path} with lr: {optimizer.param_groups[0]['lr']}")
  890. return ckpt['epoch'] + 1
  891. return 0
  893. def train_epoch(epoch, trainLoader, train_config, start_ep):
  894. student_model = train_config['student_model']
  895. teacher_model = train_config['teacher_model']
  896. optimizer = train_config['optimizer']
  897. scaler = train_config['scaler']
  898. accumulation_steps = train_config['accumulation_steps']
  899. teacher_model_paths = train_config['teacher_model_paths']
  900. task = train_config['task']
  902. student_model.train()
  903. teacher_model.eval()
  905. epoch_losses = {'total': 0, 'seg': 0, 'bce': 0}
  906. tumors_losses = {'GLI': [], 'PED': [], 'SSA': [], 'MEN': [], 'MET': []}
  908. with tqdm(total=len(trainLoader), desc=f"(Epoch {epoch + 1}/{start_ep + train_config['epochs']})", unit='batch') as pbar:
  909. optimizer.zero_grad()
  911. for step, y in enumerate(trainLoader):
  912. batch_loss = 0
  913. for sub_step, data_type in enumerate(y['data_type']):
  914. imgs = y['imgs'][sub_step].unsqueeze(0).to('cuda:0')
  915. masks = y['masks'][sub_step].unsqueeze(0).to('cuda:0')
  917. load_teacher_model(teacher_model, data_type, teacher_model_paths)
  919. with amp.autocast('cuda:0'):
  920. teacher_outputs = teacher_model(imgs)
  922. detached_teacher_output = {k: v.detach().to('cuda:1') for k, v in teacher_outputs.items()}
  923. imgs, masks = imgs.to('cuda:1'), masks.to('cuda:1')
  925. with amp.autocast('cuda:1'):
  926. student_outputs = student_model(detached_teacher_output, {'imgs': imgs, 'masks': masks})
  927. loss = (student_outputs['batch_total_student_loss'] / accumulation_steps)
  928. batch_loss += loss.item()
  929. tumors_losses[data_type].append(loss.item())
  931. scaler.scale(loss).backward()
  933. task.get_logger().report_scalar(
  934. title=f"Tumors training losses per epoch {epoch+1}",
  935. series=f"{data_type} loss",
  936. iteration=len(tumors_losses[data_type]),
  937. value=float(loss.item())
  938. )
  940. for key in epoch_losses:
  941. if key != 'total':
  942. epoch_losses[key] += (student_outputs.get(f'{key}_weighted', 0) / accumulation_steps)
  944. if (sub_step + 1) % accumulation_steps == 0 or (sub_step + 1) == len(y['data_type']):
  945. scaler.step(optimizer)
  946. scaler.update()
  947. optimizer.zero_grad()
  949. epoch_losses['total'] += batch_loss
  950. pbar.update(1)
  952. for key in epoch_losses:
  953. epoch_losses[key] /= len(trainLoader)
  955. return epoch_losses, tumors_losses
  957. def log_KD_losses_over_epochs(epoch, epoch_losses, tumors_losses, task):
  958. for loss_type, val in epoch_losses.items():
  959. task.get_logger().report_scalar(
  960. title="KD Losses over Epochs",
  961. series=f"{loss_type} loss",
  962. iteration=epoch + 1,
  963. value=epoch_losses[loss_type]
  964. )
  966. for tumor_type, losses in tumors_losses.items():
  967. task.get_logger().report_scalar(
  968. title=f"{tumor_type} Losses over Epochs",
  969. series=f"{tumor_type} Epoch trainLoss",
  970. iteration=epoch + 1,
  971. value=sum(losses) / len(losses) if losses else 0
  972. )
  974. def validate_and_save(epoch, valLoader, train_config, epoch_losses):
  975. student_model = train_config['student_model']
  976. scheduler = train_config['scheduler']
  977. optimizer = train_config['optimizer']
  978. scaler = train_config['scaler']
  979. out_checkpoint_dir = train_config['out_checkpoint_dir']
  980. task = train_config['task']
  982. val_loss = evaluate(student_model, valLoader, epoch, task)
  983. scheduler.step(val_loss)
  985. task.get_logger().report_scalar("LR", "learning_rate", iteration=epoch+1, value=optimizer.param_groups[0]['lr'])
  986. print(f"Learning rate after epoch {epoch + 1}: {optimizer.param_groups[0]['lr']}")
  988. state = {
  989. 'epoch': epoch,
  990. 'student_model': student_model.student.state_dict(),
  991. 'optimizer_student': optimizer.state_dict(),
  992. 'lr': optimizer.param_groups[0]['lr'],
  993. 'grad_scaler_state': scaler.state_dict(),
  994. 'scheduler_state_dict': scheduler.state_dict(),
  995. 'val_dice_loss': val_loss
  996. }
  998. checkpoint_path = out_checkpoint_dir / f'Student_model_after_epoch_{epoch + 1}_trainLoss_{epoch_losses["total"]:.4f}_valLoss_{val_loss:.4f}.pth'
  999. torch.save(state, checkpoint_path)
  1000. print(f"Model saved after epoch {epoch + 1}")
  1002. def run_KD(trainLoader, valLoader, args):
  1003. setup_environment(args)
  1004. teacher_model, student_model = initialize_models()
  1005. optimizer, scheduler = initialize_optimizer_scheduler(student_model, args)
  1006. scaler = amp.GradScaler('cuda:1')
  1008. teacher_model_paths = {
  1009. 'GLI': '/kaggle/input/gliomateachernewlabels/Teacher_model_after_epoch_100_trainLoss_0.5972_valLoss_0.3019.pth',
  1010. 'SSA': '/kaggle/input/africanewlabels/Teacher_model_after_epoch_67_trainLoss_1.1080_valLoss_0.5561.pth',
  1011. 'PED': '/kaggle/input/pednewlabel/Teacher_model_after_epoch_99_trainLoss_1.4512_valLoss_1.0042.pth',
  1012. 'MEN': '/kaggle/input/meningiomateachernewlabels/Teacher_model_after_epoch_85_trainLoss_0.5824_valLoss_0.3318.pth',
  1013. 'MET': '/kaggle/input/met-teacher-new-labels/Teacher_model_after_epoch_100_trainLoss_1.6278_valLoss_0.7199.pth'
  1014. }
  1016. start_epoch = load_student_checkpoint(student_model, optimizer, scaler, scheduler, args)
  1017. task = Task.init(project_name="Fairness KD 5 Tumors Models", task_name=f"ABL Study Fairness KD 5 Tumors with BCE+SEG NO KL", reuse_last_task_id=True)
  1018. task.connect(args)
  1019. task.add_tags(['ABL Study', 'BCE+SEG without KL', "Ahmed Pro"])
  1021. print(f'''Starting Knowledge Distillation:
  1022. Epochs: From {start_epoch + 1} to {start_epoch + args['epochs']}
  1023. Batch size: 5 (effective through gradient accumulation)
  1024. Learning rate: {args['learning_rate']}
  1025. Training data coming from: {args['data_dirs']}
  1026. ''')
  1028. train_config = {
  1029. 'teacher_model': teacher_model,
  1030. 'student_model': student_model,
  1031. 'optimizer': optimizer,
  1032. 'scheduler': scheduler,
  1033. 'scaler': scaler,
  1034. 'accumulation_steps': 5,
  1035. 'teacher_model_paths': teacher_model_paths,
  1036. 'out_checkpoint_dir': args['out_checkpoint_dir'],
  1037. 'task': task,
  1038. 'epochs': args['epochs']
  1039. }
  1041. for epoch in range(start_epoch, start_epoch + args['epochs']):
  1042. epoch_losses, tumors_losses = train_epoch(epoch, trainLoader, train_config, start_epoch)
  1043. log_KD_losses_over_epochs(epoch, epoch_losses, tumors_losses, task)
  1044. validate_and_save(epoch, valLoader, train_config, epoch_losses)
  1046. print("Training completed.")
  1047. task.close()
  1049. # %%
  1050. args = {
  1051. 'workers': 2,
  1052. 'epochs': 3,
  1053. 'train_batch_size': 5,
  1054. 'val_batch_size': 2,
  1055. 'test_batch_size': 1,
  1056. 'learning_rate': 1e-3,
  1057. 'weight_decay': 1e-5,
  1058. 'lambd': 0.0051,
  1059. 'data_dirs': ["/kaggle/input/bratsglioma/Training/", "/kaggle/input/bratsafrica24/", "/kaggle/input/bratsped/Training/", "/kaggle/input/bratsmen/", "/kaggle/input/bratsmet24/"],
  1060. 'in_checkpoint_dir': Path('/kaggle/input/data-abl-study-fairness-5-tumors-bce-seg-no-kl/'),
  1061. 'out_checkpoint_dir': Path('/kaggle/working/')
  1062. }
  1064. trainLoader, valLoader, testLoader = prepare_data_loaders(args)
  1065. run_KD(trainLoader, valLoader, args)
  1067. # %% [markdown]
  1068. # # Press here

ABL_(BCE+SEG).ipynb at commit 248789f, no license · at the source

Overview

Authors: Ahmed Elzayat1, Nourhan Hanafy1, Mariem Magdy1, Hazem Zakaria1, Mina Tayeh1, Abdulkhalek Al-Fakih2, Yeong Hyeon Gu2, Mohammed A. Al-masni2, Meena M. Makary1,3
  1. Systems and Biomedical Engineering Department, Faculty of Engineering, Cairo University,Cairo, Egypt
  2. Department of Artificial Intelligence and Data Science, College of Artificial Intelligence Convergence, Sejong University,Seoul, Republic of Korea
  3. Department of Radiology, Massachusetts General Hospital, Athinoula A. Martinos Center for Biomedical Imaging, Harvard Medical School,Charlestown, MA USA
Institutions: Cairo University (Egypt); Sejong University (South Korea); Harvard University (United States)
Journal: Scientific reports, volume 16, issue 1, article 12969
Dates: received 15 July 2025; accepted 7 January 2026; published online 10 March 2026
Type: Research article · Language: English
License: CC BY-NC-ND
Identifiers: DOI 10.1038/s41598-026-35627-x · PMID 41807430 · PMCID PMC13096204 · OpenAlex W7134947088
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: structural MRI / diffusion (modality), human (organism), other condition (population), methods / tools (subfield)
Methods: Connectivity, Machine learning
Keywords: Generalizability, Foundation model, Knowledge distillation, Brain tumor segmentation, Cancer, Computational biology and bioinformatics, Mathematics and computing, Neuroscience, Oncology
MeSH: Brain Neoplasms*, Glioma*, Image Processing, Computer-Assisted*, Humans, Magnetic Resonance Imaging (* major topic)
Topic: Advanced Neural Network Applications (Computer Vision and Pattern Recognition, Computer Science), according to OpenAlex
Funding: IITP (Institute of Information & Communications Technology Planning & Evaluation) - ITRC (Information Technology Research Center) grant funded by the Korea government (Ministry of Science and ICT) (IITP-2025-RS-2024-00437191); National Research Foundation of Korea (NRF) funded by the Korean government (MSIT) (No. RS-2023-00243034 and No. RS-2026-25487511)
Citations: not cited yet (Europe PMC); 53 references in the paper

Abstract

The abstract is not reproduced here: the paper's license (CC BY-NC-ND) does not allow it. Read it in the paper, at the publisher or on Europe PMC.

Repository

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

AhmeddEmad7/Brain-Tumor-Segmentation-Advancing-Generalizability

License: none: the authors keep all their rights
State: the link answers, verified on 30 September 2026
Evidence: files inventoried
Commit: 248789f547fba152cc4f61c1f4026871246c762a, 15 April 2026
Languages: Python (100), TypeScript (85), Jupyter (25), JavaScript (8), Shell (2)
Size: 482 files, 220 scripts
Software Heritage: not archived
Found in: “Data availability”
Holds: README, environment (AI/requirements.txt, software/docker-compose.yaml, software/backend/requirements.txt, software/frontend/Dockerfile, software/backend/ApiGateway/Dockerfile, software/backend/ApiGateway/requirements.txt, software/backend/Services/NiftiStorage/Dockerfile, software/backend/Services/NiftiStorage/requirements.txt, software/backend/Services/Orthanc/Dockerfile, software/backend/Services/Redis/dockerfile, software/backend/Services/Reporting/Dockerfile, software/backend/Services/Reporting/requirements.txt), documentation, 25 notebooks
Not found: license file, CITATION.cff, tests, continuous integration
Tools: PyTorch (15 files), MONAI (9 files), NumPy (8 files), NiBabel (7 files), Matplotlib (4 files)
Availability: 1 check, the latest on 30 September 2026: the link answers
  • 30 September 2026: the link answers
16 files
At the source: github.com/AhmeddEmad7/Brain-Tumor-Segmentation-Advancing-Generalizability

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;
  • 16 scripts, each with its path and the digest of its content;
  • 17 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

Code and data availability statement

The paper has a code and data availability statement. Its license (CC BY-NC-ND) does not allow reproducing it here; in short, from what the harvester recognized in it:

Read it in the paper: doi.org/10.1038/s41598-026-35627-x.

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, 9 authors, 9 keywords, 5 MeSH terms, 2 funders, 18 references.

Cite

This paper

Elzayat, A., Hanafy, N., Magdy, M., Zakaria, H., Tayeh, M., Al-Fakih, A., Gu, Y. H., Al-masni, M. A., & Makary, M. M. (2026). Distilling population specific expertise into a unified model for generalizable brain tumor segmentation. Scientific reports, 16(1), 12969. https://doi.org/10.1038/s41598-026-35627-x

BibTeX

@article{elzayat2026distilling,
author = {Elzayat, Ahmed and Hanafy, Nourhan and Magdy, Mariem and Zakaria, Hazem and Tayeh, Mina and Al-Fakih, Abdulkhalek and Gu, Yeong Hyeon and Al-masni, Mohammed A. and Makary, Meena M.},
title = {{Distilling population specific expertise into a unified model for generalizable brain tumor segmentation}},
journal = {Scientific reports},
year = {2026},
month = mar,
volume = {16},
number = {1},
pages = {12969},
publisher = {Nature Publishing Group},
issn = {2045-2322},
doi = {10.1038/s41598-026-35627-x},
url = {https://doi.org/10.1038/s41598-026-35627-x},
pmid = {41807430},
pmcid = {PMC13096204}
}

RIS

TY - JOUR
AU - Elzayat, Ahmed
AU - Hanafy, Nourhan
AU - Magdy, Mariem
AU - Zakaria, Hazem
AU - Tayeh, Mina
AU - Al-Fakih, Abdulkhalek
AU - Gu, Yeong Hyeon
AU - Al-masni, Mohammed A.
AU - Makary, Meena M.
TI - Distilling population specific expertise into a unified model for generalizable brain tumor segmentation
T2 - Scientific reports
J2 - Sci Rep
PY - 2026
DA - 2026/03/10
VL - 16
IS - 1
SP - 12969
SN - 2045-2322
PB - Nature Publishing Group
DO - 10.1038/s41598-026-35627-x
UR - https://doi.org/10.1038/s41598-026-35627-x
LA - en
ER -

CSL-JSON

{
"id": "10.1038/s41598-026-35627-x",
"type": "article-journal",
"title": "Distilling population specific expertise into a unified model for generalizable brain tumor segmentation",
"container-title": "Scientific reports",
"author": [
{
"family": "Elzayat",
"given": "Ahmed"
},
{
"family": "Hanafy",
"given": "Nourhan"
},
{
"family": "Magdy",
"given": "Mariem"
},
{
"family": "Zakaria",
"given": "Hazem"
},
{
"family": "Tayeh",
"given": "Mina"
},
{
"family": "Al-Fakih",
"given": "Abdulkhalek"
},
{
"family": "Gu",
"given": "Yeong Hyeon"
},
{
"family": "Al-masni",
"given": "Mohammed A."
},
{
"family": "Makary",
"given": "Meena M."
}
],
"container-title-short": "Sci Rep",
"volume": "16",
"issue": "1",
"page": "12969",
"DOI": "10.1038/s41598-026-35627-x",
"PMID": "41807430",
"PMCID": "PMC13096204",
"ISSN": "2045-2322",
"publisher": "Nature Publishing Group",
"URL": "https://doi.org/10.1038/s41598-026-35627-x",
"language": "en",
"issued": {
"date-parts": [
[
2026,
3,
10
]
]
}
}

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.1016/j.patter.2026.101538 [code]
A multi-modal foundation model for brain disease diagnosis and medical imaging.
Journal: Patterns (New York, N.Y.)
In common: MONAI, NiBabel, PyTorch, 2 other tools, 3 references
[2] doi:10.1371/journal.pone.0354511 [code]
TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging segmentation.
Journal: PloS one
In common: MONAI, NiBabel, PyTorch, 2 other tools, methods / tools, other condition, 2 references
[3] doi:10.1038/s41598-026-48496-1 [code]
A unified FLAIR hyperintensity segmentation model for various CNS tumor types and acquisition time points.
Journal: Scientific reports
In common: NiBabel, Matplotlib, NumPy, methods / tools, structural MRI / diffusion, other condition, 4 references
[4] 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, Matplotlib, 1 other tool, methods / tools, structural MRI / diffusion, other condition, 3 references
[5] doi:10.1158/2767-9764.crc-25-0710 [code]
MRI Deep Learning for Differentiating Glioblastoma, IDH Wild-type from Central Nervous System Diffuse Large B-cell Lymphoma.
Journal: Cancer research communications
In common: MONAI, NiBabel, PyTorch, 2 other tools, structural MRI / diffusion, other condition, 1 reference
[6] doi:10.1371/journal.pdig.0001316 [code]
Reliability of a convolutional neural network in segmenting multiple sclerosis lesions from MRI: Impact of data augmentation, image modality and tolerance with U-Net architecture.
Journal: PLOS digital health
In common: MONAI, NiBabel, PyTorch, 2 other tools, methods / tools, structural MRI / diffusion, 1 reference
[7] doi:10.3390/diagnostics16172806
Brain Tumor Segmentation and Grading on MRI Using Deep Learning: A Systematic Literature Review and Benchmark-Driven Comparative Analysis.
Journal: Diagnostics (Basel, Switzerland)
In common: methods / tools, structural MRI / diffusion, other condition, 5 references
[8] doi:10.1038/s41598-026-54446-8 [code]
Deep learning-based Desikan-Killiany parcellation of the brain using diffusion MRI.
Journal: Scientific reports
In common: MONAI, NiBabel, PyTorch, 1 other tool, methods / tools, structural MRI / diffusion, 1 reference
[9] doi:10.3389/fnins.2026.1870124 [code]
An end-to-end pipeline for automated fetal brain segmentation and biometry from 3D SSFP MRI.
Journal: Frontiers in neuroscience
In common: MONAI, NiBabel, PyTorch, 2 other tools, structural MRI / diffusion, 1 reference
[10] doi:10.1371/journal.pone.0351405 [code]
Unimodal vs. multimodal deep learning for non-invasive MGMT promoter methylation prediction in glioblastoma: A systematic evaluation on the BraTS 2021 dataset.
Journal: PloS one
In common: PyTorch, NumPy, methods / tools, structural MRI / diffusion, 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.

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: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

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