Skip to content

README.md

Latest commit

 

History

History
129 lines (92 loc) · 7.65 KB

README.md

File metadata and controls

129 lines (92 loc) · 7.65 KB
Raw

mriaug

pip install mriaug to use a 3D image library that is ~50x faster and simpler than torchio by

  • only using PyTorch → full GPU(+autograd) support 🔥
  • being tiny: ~200 lines of code → no room for bugs 🐛

while offering ~20 different augmentations (incl. MRI-specific operations) 🩻

👶 Normal users should use mriaug via niftiai, a deep learning framework for 3D images, since it

  • provides aug_transforms3d: A convenient function that compiles all mriaugmentations!
  • simplifies all the code needed for data loading, training, visualization...check it out here!

👴 Experienced users can build their own framework upon mriaug (use niftiai/augment.py as a cheat sheet)

Usage 💡

Let's create a 3D image tensor (with additional batch and channel dimension) and apply flip3d

import torch
from mriaug import flip3d

shape = (1, 1, 4, 4, 4)
x = torch.linspace(0, 1, 4**3).view(*shape)
x_flipped = flip3d(x)
print(x[..., 0, 0])  # tensor([[[0.0000, 0.2540, 0.5079, 0.7619]]])
print(x_flipped[..., 0, 0])  # tensor([[[0.7619, 0.5079, 0.2540, 0.0000]]])

Explore the gallery to understand the usage and effect of all ~20 augmentations!

Speed 💨

The popular libraries torchio and MONAI (utilizes torchio) often use ITK (CPU only) like this

PyTorch tensor → NumPy array → NiBabel image → ITK operation (C/C++) → NumPy array → PyTorch tensor

to augment a PyTorch tensor 😬 That's complicated and does not use the (for neural net training needed) GPU 🐌

Instead, mriaug directly uses PyTorch (CPU & GPU support) resulting in

  • ~50x fewer lines of code: torchio: ~10,000 LOC, mriaug: ~200 LOC 🤓
  • ~50x speedup on GPU 🔥 based on the table below (run speed.py to reproduce) 💨
Click here, to see runtimes on a 256³ image in seconds (on AMD 5950X CPU and NVIDIA RTX 3090 GPU)
Transformation torchio mriaug on CPU mriaug on GPU Speedup vs. torchio
Flip 0.014 0.012 0.002 7.5x
Affine 0.297 0.608 0.011 27.9x
Warp 0.951 0.850 0.009 103.3x
Bias Field 3.258 0.081 0.002 1813.0x
Noise 0.117 0.105 0.001 230.4x
Downsample 0.282 0.013 0.000 592.3x
Ghosting 0.241 0.170 0.003 78.3x
Spike 0.265 0.172 0.003 88.8x
Motion 0.696 0.540 0.009 78.6x

Gallery 🧠

Let's load an example 3D image x, show it with niftiview (used to create all images below)

define some arguments

size = (160, 196, 160)
zoom = torch.tensor([[-.2, 0, 0]])
rotate = torch.tensor([[0, .1, 0]])
translate = torch.tensor([[0, 0, .2]])
shear = torch.tensor([[0, .05, 0]])

and run all augmentations (see runall.py):

flip3d(x, dim=0)

dihedral3d(x, k=2)

crop3d(x, translate, size)

zoom3d(x, zoom)

rotate3d(x, rotate)

translate3d(x, translate)

shear3d(x, shear)

affine3d(x, zoom, rotate, translate, shear)

warp3d(x, magnitude=.01)

affinewarp3d(x, zoom, rotate, translate, shear, magnitude=.01)

bias_field3d(x, intensity=.2)

contrast(x, lighting=.5)

noise3d(x, intensity=.05)

chi_noise3d(x, intensity=.05, dof=3) set dof=2 for Rician noise

downsample3d(x, scale=.25, dim=2)

ghosting3d(x, intensity=.5)

spike3d(x, intensity=.2)

ringing3d(x, intensity=.5)

motion3d(x, intensity=.5)