AlignShift

NEW: Code for our new MICCAI’21 paper “Asymmetric 3D Context Fusion for Universal Lesion Detection” will also be pushed to this repository soon.

AlignShift: Bridging the Gap of Imaging Thickness in 3D Anisotropic Volumes (MICCAI’20, early accepted)

Key contributions

  • AlignShift aims at a plug-and-play replacement of standard 3D convolution for 3D medical images, which enables 2D-to-3D pretraining as ACS Convolutions. It converts theoretically any 2D pretrained network into thickness-aware 3D network.
  • AlignShift bridges the performance gap between thin- and thick-slice volumes by a unified framework. Remarkably, the AlignShift-converted networks behave like 3D for the thin-slice, nevertheless degenerate to 2D for the thick-slice adaptively.
  • Without whistles and bells, we outperform previous state of the art by considerable margins on large-scale DeepLesion benchmark for universal lesion detection.

Code structure

  • alignshift the core implementation of AlignShift convolution and TSM convolution, including the operators, models, and 2D-to-3D/AlignShift/TSM model converters.
    • operators: include AlignShiftConv, TSMConv.
    • converters.py: include converters which convert 2D models to 3dConv/AlignShiftConv/TSMConv counterparts.
    • models: Native AlignShift/TSM models.
  • deeplesion the experiment code is base on mmdetection ,this directory consists of compounents used in mmdetection.
  • mmdet

Installation

  • git clone this repository
  • pip install -e .

Convert a 2D model into 3D with a single line of code

from converter import Converter
import torchvision
from alignshift import AlignShiftConv
# m is a standard pytorch model
m = torchvision.models.resnet18(True)
alignshift_conv_cfg = dict(conv_type=AlignShiftConv, 
                          n_fold=8, 
                          alignshift=True, 
                          inplace=True,
                          ref_spacing=0.2, 
                          shift_padding_zero=True)
m = Converter(m, 
              alignshift_conv_cfg, 
              additional_forward_fts=['thickness'], 
              skip_first_conv=True, 
              first_conv_input_channles=1)
# after converted, m is using AlignShiftConv and capable of processing 3D volumes
x = torch.rand(batch_size, in_channels, D, H, W)
thickness = torch.rand(batch_size, 1)
out = m(x, thickness)

Usage of AlignShiftConv/TSMConv operators

from alignshift.operators import AlignShiftConv, TSMConv
x = torch.rand(batch_size, 3, D, H, W)
thickness = torch.rand(batch_size, 1)
# AlignShiftConv to process 3D volumnes
conv = AlignShiftConv(in_channels=3, out_channels=10, kernel_size=3, padding=1, n_fold=8, alignshift=True, ref_thickness=2.0)
out = conv(x, thickness)
# TSMConv to process 3D volumnes
conv = TSMConv(in_channels=3, out_channels=10, kernel_size=3, padding=1, n_fold=8, tsm=True)
out = conv(x)

Usage of native AlignShiftConv/TSMConv models

from alignshift.models import DenseNetCustomTrunc3dAlign, DenseNetCustomTrunc3dTSM
net = DenseNetCustomTrunc3dAlign(num_classes=3)
B, C_in, D, H, W = (1, 3, 7, 256, 256)
input_3d = torch.rand(B, C_in, D, H, W)
thickness = torch.rand(batch_size, 1)
output_3d = net(input_3d, thickness)

How to run the experiments

  • Dataset

  • Preparing mmdetection script

  • Model weights

    Our trained weights published on:

  • Training

    ./deeplesion/train_dist.sh ${mmdetection script} ${dist training GPUS}
    • Train AlignShiftConv models

    ./deeplesion/train_dist.sh ./deeplesion/mconfig/densenet_align.py 2
    • Train TSMConv models

    ./deeplesion/train_dist.sh ./deeplesion/mconfig/densenet_tsm.py 2

  • Evaluation

    ./deeplesion/eval.sh ${mmdetection script} ${checkpoint path}

    ./deeplesion/eval.sh ./deeplesion/mconfig/densenet_align.py ./deeplesion/model_weights/alignshift_7slice.pth

GitHub

https://github.com/M3DV/AlignShift