changes

.gitignore

@@ -1,19 +1,25 @@
 # Byte-compiled / optimized / DLL files
 __pycache__/
+notebooks/medical_seg/data/
 *.py[cod]
 *$py.class
 *.vscode
 __pycache__
 *.pyc
 *.idea
 *.DS_Store
-
+tmp*
+tests/integration/
+zz*
 # data which is automatically downloaded inside notebooks
 notebooks/ExBox3
+notebooks/MedNIST.tar.gz
+notebooks/MedNIST/
+notebooks/demo_classification.py
 
 # C extensions
 *.so
-
+notebooks/MedNIST
 # Distribution / packaging
 .Python
 build/

.pre-commit-config.yaml

@@ -1,5 +1,5 @@
 default_language_version:
-  python: python3.8
+  python: python3.9
 
 ci:
   autofix_prs: true
@@ -19,6 +19,7 @@ repos:
       - id: pretty-format-json
       - id: check-added-large-files
         exclude: .*\.ipynb
+        args: ['--maxkb=1000']
       - id: check-docstring-first
       - id: detect-private-key
 
@@ -35,7 +36,7 @@ repos:
       - id: yesqa
 
   - repo: https://github.com/psf/black
-    rev: 21.7b0
+    rev: 22.3.0
     hooks:
       - id: black
         name: Format code
@@ -47,7 +48,7 @@ repos:
         name: imports
 
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.7
+    rev: 0.7.10
     hooks:
       - id: mdformat
         additional_dependencies:

CITATION.cff

@@ -0,0 +1,18 @@
+cff-version: 1.2.0
+title: >-
+  PhoenixDL/rising: High-Performance Differentiable
+  Medical Data Augmentation
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Justus
+    family-names: Schock
+    affiliation: Contributed Equally
+    orcid: 'https://orcid.org/0000-0003-0512-3053'
+  - given-names: Michael
+    family-names: Baumgartner
+    affiliation: Contributed Equally
+  - given-names: Leon
+    family-names: Weninger

README.md

@@ -7,7 +7,7 @@
 ![PyPI - License](https://img.shields.io/pypi/l/rising)
 [![Chat](https://img.shields.io/badge/Slack-PhoenixDL-orange)](https://join.slack.com/t/phoenixdl/shared_invite/enQtODgwODI0MTE1MjgzLTJkZDE4N2NhM2VmNzVhYTEyMzI3NzFmMDY0NjM3MzJlZWRmMTk5ZWM1YzY2YjY5ZGQ1NWI1YmJmOTdiYTdhYTE)
 [![Documentation Status](https://readthedocs.org/projects/rising/badge/?version=latest)](https://rising.readthedocs.io/en/latest/?badge=latest)
-[![pre-commit.ci status](https://results.pre-commit.ci/badge/github/PhoenixDL/rising/master.svg)](https://results.pre-commit.ci/latest/github/PhoenixDL/rising/master)
+[![pre-commit.ci status](https://results.pre-commit.ci/badge/github/PhoenixDL/rising/master.svg)](https://results.pre-commit.ci/latest/github/PhoenixDL/rising/master)[![DOI](https://zenodo.org/badge/222287924.svg)](https://zenodo.org/badge/latestdoi/222287924)
 
 </div>
 
@@ -17,6 +17,19 @@
 | ![Python](https://img.shields.io/badge/python-3.6/3.7/3.8-orange) | ![System](https://img.shields.io/badge/Windows-blue) | ![Unittests Windows](https://github.com/PhoenixDL/rising/workflows/Unittests%20Windows/badge.svg) |
 | ![Python](https://img.shields.io/badge/python-3.6/3.7/3.8-orange) | ![System](https://img.shields.io/badge/MacOS-blue)   | ![Unittests macOS](https://github.com/PhoenixDL/rising/workflows/Unittests%20MacOS/badge.svg)     |
 
+## What I have modified from the original repo?
+
+I noticed that rising under its current form has small bugs and inconsistency in terms of design and performance, especially for medical image segmentation task with 3D volumes. Based on project requirement, several improvements have been made so that it can now enjoy a better compatibility with 3D datas and intenstive augmentation operations using GPUs. Issues and PRs are welcome.
+
+Thanks a lot @Yuxiang1990 for the kind PR.
+
+## Working example on ACDC dataset
+
+[ACDC example](notebooks/medical_seg/acdc_seg.py) should provide a resonable example of how to use rising besides the below
+explanations.
+
+______________________________________________________________________
+
 ## What is `rising`?
 
 Rising is a high-performance data loading and augmentation library for 2D *and* 3D data completely written in PyTorch.

notebooks/medical_seg/acdc_seg.py

@@ -0,0 +1,66 @@
+from tqdm import tqdm
+
+from notebooks.medical_seg.dataset import ACDCDataset, InfiniteRandomSampler
+from rising import transforms
+from rising.constants import FInterpolation
+from rising.loading import DataLoader, default_transform_call
+from rising.random import UniformParameter
+
+tra_dataset = ACDCDataset(root="/home/jizong/Workspace/rising/notebooks/medical_seg/data", train=True)
+
+seq_tra_augment = transforms.Compose(
+    transforms.NormPercentile(keys=("image",), min=0.02, max=0.98),
+    transforms.ResizeNative(
+        size=(10, 224, 224),
+        mode=(FInterpolation.trilinear, FInterpolation.nearest),
+        preserve_range=True,
+        keys=("image", "label"),
+    ),
+    # transforms.PadRandomCrop(size=(10, 224, 224), pad_size=2, pad_value=(0, 0), keys=("image", "label")),
+)
+
+batch_augment = transforms.Compose(
+    # transforms.ToDtype(keys=("image", "label"), dtype=torch.half),
+    # transforms.ToDevice(keys=("image", "label"), device=torch.device("cuda")),
+    transforms.GammaCorrection(gamma=UniformParameter(0.8, 2), keys=("image",)),
+    transforms.RicianNoiseTransform(keys=("image",), std=0.05, keep_range=False),
+    transforms.BaseAffine(
+        scale=(1, UniformParameter(0.5, 4), UniformParameter(0.9, 1.2)),
+        rotation=(0, UniformParameter(-10, 10), UniformParameter(-10, 10)),
+        degree=True,
+        p=1,
+        per_sample=True,
+        interpolation_mode=("bilinear", "nearest"),
+        keys=("image", "label"),
+    ),
+    transforms.RandomCrop(size=(8, 192, 168), keys=("image", "label")),
+    transforms.ElasticDistortion(
+        std=20,
+        alpha=0.2,
+        dim=3,
+        keys=("image", "label"),
+        interpolation_mode=("bilinear", "nearest"),
+    ),
+    transform_call=default_transform_call,
+)
+
+tra_loader = DataLoader(
+    tra_dataset,
+    sampler=InfiniteRandomSampler(tra_dataset, shuffle=False),
+    batch_size=6,
+    sample_transforms=seq_tra_augment,
+    gpu_transforms=batch_augment,
+    pseudo_batch_dim=True,
+    num_workers=16,
+)
+
+for data in tqdm(tra_loader):
+    image, label = data["image"], data["label"]
+    from tests.realtime_viewer import multi_slice_viewer_debug
+
+    multi_slice_viewer_debug([*image.squeeze()], *label.squeeze(), block=True, no_contour=True)
+    # from tests.realtime_viewer import multi_slice_viewer_debug
+    #
+    # for img, lab in zip(image, label):
+    #     multi_slice_viewer_debug(img.squeeze(), lab.squeeze(), block=False, no_contour=True)
+    # plt.show()

notebooks/medical_seg/data/train/patient001/Info.cfg

@@ -0,0 +1,6 @@
+ED: 1
+ES: 12
+Group: DCM
+Height: 184.0
+NbFrame: 30
+Weight: 95.0

notebooks/medical_seg/dataset.py

@@ -0,0 +1,72 @@
+import re
+from collections import Iterator
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+from torch.utils.data import Sampler
+from torch.utils.data.dataset import Dataset, T_co
+
+
+class ACDCDataset(Dataset):
+    def __init__(self, *, root: str, train: bool = True) -> None:
+        self.root = root
+        self.train = train
+        self._root = Path(root, "train" if train else "val")
+
+        self.images = filter(
+            self.image_filter, [str(x.relative_to(self._root)) for x in Path(self._root).rglob("*.nii.gz")]
+        )
+        self.images = sorted(self.images)
+
+    def __getitem__(self, index) -> T_co:
+        image_path = str(self._root / self.images[index])
+        gt_path = image_path.replace(".nii.gz", "_gt.nii.gz")
+        image = sitk.GetArrayFromImage(sitk.ReadImage(image_path)).astype(float, copy=False)[None, ...]
+        gt = sitk.GetArrayFromImage(sitk.ReadImage(gt_path)).astype(float, copy=False)[None, ...]
+        return {"image": torch.from_numpy(image), "label": torch.from_numpy(gt)}
+
+    def __len__(self):
+        return len(self.images)
+
+    @staticmethod
+    def image_filter(path: str):
+        _match = re.compile(r"patient\d+_frame\d+.nii.gz").search(str(path))
+        if _match is None:
+            return False
+        return True
+
+
+class _InfiniteRandomIterator(Iterator):
+    def __init__(self, data_source, shuffle=True):
+        self.data_source = data_source
+        self.shuffle = shuffle
+        if self.shuffle:
+            self.iterator = iter(torch.randperm(len(self.data_source)).tolist())
+        else:
+            self.iterator = iter(torch.arange(start=0, end=len(self.data_source)).tolist())
+
+    def __next__(self):
+        try:
+            idx = next(self.iterator)
+        except StopIteration:
+            if self.shuffle:
+                self.iterator = iter(torch.randperm(len(self.data_source)).tolist())
+            else:
+                self.iterator = iter(torch.arange(start=0, end=len(self.data_source)).tolist())
+            idx = next(self.iterator)
+        return idx
+
+
+class InfiniteRandomSampler(Sampler):
+    def __init__(self, data_source, shuffle=True):
+        super().__init__(data_source)
+        self.data_source = data_source
+        self.shuffle = shuffle
+
+    def __iter__(self):
+        return _InfiniteRandomIterator(self.data_source, shuffle=self.shuffle)
+
+    def __len__(self):
+        return len(self.data_source)

notebooks/medical_seg/unet_3d.py

@@ -0,0 +1,130 @@
+# 3D-UNet model.
+# x: 128x128 resolution for 32 frames.
+import torch
+import torch.nn as nn
+
+
+def conv_block_3d(in_dim, out_dim, activation):
+    return nn.Sequential(
+        nn.Conv3d(in_dim, out_dim, kernel_size=3, stride=1, padding=1),
+        nn.BatchNorm3d(out_dim),
+        activation,
+    )
+
+
+def conv_trans_block_3d(in_dim, out_dim, activation):
+    return nn.Sequential(
+        nn.ConvTranspose3d(in_dim, out_dim, kernel_size=3, stride=2, padding=1, output_padding=1),
+        nn.BatchNorm3d(out_dim),
+        activation,
+    )
+
+
+def max_pooling_3d():
+    return nn.MaxPool3d(kernel_size=2, stride=2, padding=0)
+
+
+def conv_block_2_3d(in_dim, out_dim, activation, kernel_size=(3, 3, 3)):
+    return nn.Sequential(
+        conv_block_3d(in_dim, out_dim, activation),
+        nn.Conv3d(out_dim, out_dim, kernel_size=kernel_size, stride=1, padding=1),
+        nn.BatchNorm3d(out_dim),
+    )
+
+
+class UNet(nn.Module):
+    def __init__(self, in_dim, out_dim, num_filters):
+        super(UNet, self).__init__()
+
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_filters = num_filters
+        activation = nn.LeakyReLU(0.2, inplace=True)
+
+        # Down sampling
+        self.down_1 = conv_block_2_3d(self.in_dim, self.num_filters, activation, kernel_size=(1, 3, 3))
+        self.pool_1 = max_pooling_3d()
+        self.down_2 = conv_block_2_3d(self.num_filters, self.num_filters * 2, activation, kernel_size=(1, 3, 3))
+        self.pool_2 = max_pooling_3d()
+        self.down_3 = conv_block_2_3d(self.num_filters * 2, self.num_filters * 4, activation, kernel_size=(1, 3, 3))
+        self.pool_3 = max_pooling_3d()
+        self.down_4 = conv_block_2_3d(self.num_filters * 4, self.num_filters * 8, activation)
+        self.pool_4 = max_pooling_3d()
+        self.down_5 = conv_block_2_3d(self.num_filters * 8, self.num_filters * 16, activation)
+        self.pool_5 = max_pooling_3d()
+
+        # Bridge
+        self.bridge = conv_block_2_3d(self.num_filters * 16, self.num_filters * 32, activation)
+
+        # Up sampling
+        self.trans_1 = conv_trans_block_3d(self.num_filters * 32, self.num_filters * 32, activation)
+        self.up_1 = conv_block_2_3d(self.num_filters * 48, self.num_filters * 16, activation)
+        self.trans_2 = conv_trans_block_3d(self.num_filters * 16, self.num_filters * 16, activation)
+        self.up_2 = conv_block_2_3d(self.num_filters * 24, self.num_filters * 8, activation)
+        self.trans_3 = conv_trans_block_3d(self.num_filters * 8, self.num_filters * 8, activation)
+        self.up_3 = conv_block_2_3d(self.num_filters * 12, self.num_filters * 4, activation)
+        self.trans_4 = conv_trans_block_3d(self.num_filters * 4, self.num_filters * 4, activation)
+        self.up_4 = conv_block_2_3d(self.num_filters * 6, self.num_filters * 2, activation)
+        self.trans_5 = conv_trans_block_3d(self.num_filters * 2, self.num_filters * 2, activation)
+        self.up_5 = conv_block_2_3d(self.num_filters * 3, self.num_filters * 1, activation)
+
+        # Output
+        self.out = conv_block_3d(self.num_filters, out_dim, activation)
+
+    def forward(self, x):
+        # Down sampling
+        down_1 = self.down_1(x)  # -> [1, 4, 128, 128, 128]
+        pool_1 = self.pool_1(down_1)  # -> [1, 4, 64, 64, 64]
+
+        down_2 = self.down_2(pool_1)  # -> [1, 8, 64, 64, 64]
+        pool_2 = self.pool_2(down_2)  # -> [1, 8, 32, 32, 32]
+
+        down_3 = self.down_3(pool_2)  # -> [1, 16, 32, 32, 32]
+        pool_3 = self.pool_3(down_3)  # -> [1, 16, 16, 16, 16]
+
+        down_4 = self.down_4(pool_3)  # -> [1, 32, 16, 16, 16]
+        pool_4 = self.pool_4(down_4)  # -> [1, 32, 8, 8, 8]
+
+        down_5 = self.down_5(pool_4)  # -> [1, 64, 8, 8, 8]
+        pool_5 = self.pool_5(down_5)  # -> [1, 64, 4, 4, 4]
+
+        # Bridge
+        bridge = self.bridge(pool_5)  # -> [1, 128, 4, 4, 4]
+
+        # Up sampling
+        trans_1 = self.trans_1(bridge)  # -> [1, 128, 8, 8, 8]
+        concat_1 = torch.cat([trans_1, down_5], dim=1)  # -> [1, 192, 8, 8, 8]
+        up_1 = self.up_1(concat_1)  # -> [1, 64, 8, 8, 8]
+
+        trans_2 = self.trans_2(up_1)  # -> [1, 64, 16, 16, 16]
+        concat_2 = torch.cat([trans_2, down_4], dim=1)  # -> [1, 96, 16, 16, 16]
+        up_2 = self.up_2(concat_2)  # -> [1, 32, 16, 16, 16]
+
+        trans_3 = self.trans_3(up_2)  # -> [1, 32, 32, 32, 32]
+        concat_3 = torch.cat([trans_3, down_3], dim=1)  # -> [1, 48, 32, 32, 32]
+        up_3 = self.up_3(concat_3)  # -> [1, 16, 32, 32, 32]
+
+        trans_4 = self.trans_4(up_3)  # -> [1, 16, 64, 64, 64]
+        concat_4 = torch.cat([trans_4, down_2], dim=1)  # -> [1, 24, 64, 64, 64]
+        up_4 = self.up_4(concat_4)  # -> [1, 8, 64, 64, 64]
+
+        trans_5 = self.trans_5(up_4)  # -> [1, 8, 128, 128, 128]
+        concat_5 = torch.cat([trans_5, down_1], dim=1)  # -> [1, 12, 128, 128, 128]
+        up_5 = self.up_5(concat_5)  # -> [1, 4, 128, 128, 128]
+
+        # Output
+        out = self.out(up_5)  # -> [1, 3, 128, 128, 128]
+        return out
+
+
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    image_size = 32
+    x = torch.Tensor(2, 3, image_size, image_size, image_size)
+    x = x.to(device)
+    print("x size: {}".format(x.size()))
+
+    model = UNet(in_dim=3, out_dim=3, num_filters=16).to(device)
+
+    out = model(x)
+    print("out size: {}".format(out.size()))

requirements/install.txt

@@ -2,3 +2,4 @@ numpy
 torch>=1.6 # before 1.6 torch.searchsorted is not present
 threadpoolctl
 tqdm
+SimpleITK