Welcome to Patchwork CNN Toolbox

This toolbox comprises a framework for stacking patch-based CNNs for the integration of global information, which is beyond the size of the patch. The idea is to forward information from large, coarsly sampled patches to smaller oness that are contained in the previous patch. The Figure below shows the principle.

The toolbox builds on TensorFlow (>2.0), Keras and NiBabel. It is mainly meant for biomedical image segmentation, i.e. there is a simple interface to load data from Nifti repositories. The toolbox implements 2D and 3D patching and several data augmentation possibilities.

Installation

Either install via pip

pip install git+https://reisert@bitbucket.org/reisert/patchwork.git

or clone the repository

$ git clone https://reisert@bitbucket.org/reisert/patchwork.git

and install the requirements manually

#!python
pip install tensorflow==2 matplotlib  nibabel

In the following we explain the basic usage of the toolbox. There are also notebook on colab, or here an example for the Hippocampus task on the Medical Imaging Decathlon

Specifying your patchwork model

A patchwork CNN consists of a definition of the patching strategy (CropGenerator) and a network definition (PatchWorkModel). The patching strategy contains information about how many patch levels will be used (depth), the size and the resolution of the patches in each level and what kind of interpolation technique is used (nearest neighbor or linear). The model defines the constructors of the level specific CNNs, activations etc. .

Patching strategy

Here an example call for the CropGenerator in 2D. Note that all 3D pendants exist (ndim,nD = 2,3).

#!python

cgen = CropGenerator(ndim=2,
                  scheme = { 
                          "patch_size":[64,32],                
                          "destvox_mm": None,
                          "destvox_rel":[3,2],
                          "fov_mm":None,
                          "fov_rel":[0.75,0.6],
                            },
                  smoothfac_data = 0,
                  smoothfac_label = 0,
                  interp_type = 'NN', 
                  scatter_type = 'NN',
                  depth=4)

and here a short description of the involved parameters

scheme concludes the parameter of the way patches are drawn
patch_size specifies the pixel/voxel dimensions of the path by a tuple (2D) or a triple (3D), or a list of tuples/triples if the patch sizes are dependend on depth.
destvox_mm or destvox_rel are mutual exclusive and determine the voxel sizes of the output of the final network reponse, either in millimeter (mm) or relative to input voxel sizes.
fov_mm and fov_rel determine the field of view of the network, i.e. the size of patch in the coarsest level, either in millimiter (mm) or relative to the total size of the input.
interp_type and scatter_type defines how interpolation of patches and scattering of results are accomplishes. Possible values are nearest Neighbor ('NN') or linear interpolation ('lin').
The number of patch levels is controlled by depth >= 1

To see what the patching strategy from above is doing you can execute the following code for generating patching.

#!python

cgen.sample(tf.ones([1,372,208,5]),None,generate_type='random',
                                    resolutions=[1,3],
                                    num_patches=1,
                                    verbose=True)


>> output
input:  shape:[372,208]  width(mm):[372.0,624.0]  voxsize:[1.00,3.00]
level 0:  shape:[64,32]  width(mm):[279.0,374.4]  voxsize:[4.36,11.70]   (rel. to input:[4.36,3.90])
  dest_shape:[85,53]
level 1:  shape:[64,32]  width(mm):[246.3,299.7]  voxsize:[3.85,9.36]   (rel. to input:[3.85,3.12])
  dest_shape:[96,66]
level 2:  shape:[64,32]  width(mm):[217.5,239.9]  voxsize:[3.40,7.50]   (rel. to input:[3.40,2.50])
  dest_shape:[109,83]
level 3:  shape:[64,32]  width(mm):[192.0,192.0]  voxsize:[3.00,6.00]   (rel. to input:[3.00,2.00])
  dest_shape:[124,104]

In case of generate_type='tree' the patches are distributed in an n-tree fashion to obtain a full coverage.

The Patchwork model

The patchwork model needs a CropGenerator and function creating a keras layer, which is holding the CNN applied on each patch level. The function is expected to take a level argument and the number of outputs. Similarly, there might also be a preprocCreator and finalBlock. In case you want to have a loss also at the upper patch level and not only for the final level, turn intermediate_loss. Note that, then the loss function has to be an array of length cgen.depth. The output feature dimensions of the intermediate path layers is always num_labels+intermediate_out.

#!python


model = PatchWorkModel(cgen, 
                      lambda level,outK : createBlock(level=level,outK=outK),
                      preprocCreator = None,
                      intermediate_loss=False,
                      intermediate_out=5,
                      finalBlock=layers.Activation('sigmoid'),                      
                      num_labels = 2,
                      modelname='atest'
                      )

The function createBlock should return a keras layer of your preference. You could just take one of predefined U-net like networks:

#!python

createBlock = lambda l,K: createUnet_v2(depth=4,outK=K,feature_dim=[4,8,16,32],nD=2)

Loading your data

Here's an example of some Python code:

#!python


# loads nifti data into tf.tensors
contrasts = [ { 'subj1' :  '/content/data/brain/subj1/t1.nii' ,
                'subj2' :  '/content/data/brain/subj2/t1.nii' },
              { 'subj1' :  '/content/data/brain/subj1/flair.nii' ,
                'subj2' :  '/content/data/brain/subj2/flair.nii' } ]
labels   = [ { 'subj1' :  '/content/data/brain/subj1/PV.nii.gz' ,
                'subj2' :  '/content/data/brain/subj2/PV.nii.gz' },
              { 'subj1' :  '/content/data/brain/subj1/ML.nii.gz' ,
                'subj2' :  '/content/data/brain/subj2/ML.nii.gz' } ]

subjects = [ 'subj1','subj2' ];

data,labels,resolutions,subjs =  load_data_structured(contrasts, labels, subjects=subjects,
                           max_num_data=None,
                           nD=3,
                           ftype=tf.float32)

Training your model

A most simple way to start the training

#!python

# and now train 

model.train(data,labels,resolutions=resolutions
            epochs=5,
            num_its=100,             
            traintype='random',
            augment=None,
            num_patches=50)

Here for each element in the list data/labels 50 random patch stacks are drawn and trained for 5 epochs. This is repeated num_its=50 times.

A more complex example with augmentation and validation set.

#!python

model.train(data,labels,resolutions=resolutions
            valid_ids = [0,1,2],
            augment={'dphi':0.1,
                     'dscale':[0.1,0.1],
                     'flip':[1,0]},
            balance={'ratio':0.5,
                     'label_weight':[1]},
            num_patches=50,
            num_its=50,
            epochs=20)

The first three elements [0,1,2] of data are excluded from training and are used for validation. For augmentation patches are rotated randomly by an amount of dphi (sigmg of a normally distributed angle in radians), randomly scaled by 0.1=10% in both dimensions (dscale) and flipped along the first dimension (flip).

Applying the network

There are two different types a application modes, either random placing of patches, or a full coverage in a tree-like fashion. Additionally, there is a lazy evaluation strategy for faster execution. Here, only child patches are evaluated, if the parent gives a certain contribution.

You can either apply the network on a tensor-type image (here in the tree like fashion),

#!python
result = model.apply_full(data[0][0:1,:,:,:,:], generate_type='tree', repetitions=1, verbose=True, scale_to_original=True)

or apply it directly on a nifti with random sampling and lazy evaluation.

#!python
input = ['/content/data/brain/subj1/t1.nii' , '/content/data/brain/subj1/flair.nii' ]
result = '/content/data/brain/subj1/pred.nii'
model.apply_on_nifti(input ,result, generate_type='random',repetitions=150, 
                         branch_factor=2,
                         lazyEval={'fraction':0.5, 'reduceFun':tf.reduce_max},
                         num_chunks=5,
                         scale_to_original=False);

Auxiliaries

A CNNblock

To specify CNN architecture in a convenient way, patchwork has a builtin CNNblock, which is defined via a dictionary. But first some simple definitions for the basic layers used:

#!python

def BNrelu():
  return [layers.BatchNormalization(), layers.LeakyReLU()]
def conv(fdim):
   return layers.Conv3D(fdim,5,padding='SAME') 
def conv_up(fdim):
   return layers.Conv3DTranspose(fdim,5,padding='SAME',strides=(2,2,2))

And here an example for the most simple U-net via a CNNblock:

#!python


unet = patchwork.CNNblock({ '1001' :  [{'f': conv(10) } , {'f': conv(10), 'dest': '1004' }  ],
                  '1002' :  BNrelu() + [layers.MaxPooling3D(pool_size=(2,2,2)) ],
                  '1003' :  conv_up(10),
                  '1004' :  BNrelu(),
                  '1005' :  [layers.Dropout(rate=0.5), conv(1)] })

The semantics is as follows: the layers are always executed in a lexicographically sorted way (by dictionary key). If the value is a keras.layer it is just applied as is. If the value is a list of layers, the layers are ran sequentially (similar to keras.Sequential). If the array items are again dictionaries (like in line '1001' above), the operations of the array are NOT applied sequqntially, but concurrently, while each operation (specified by 'f') has additionally a destination layer, which is specified by 'dest'. If 'dest' is not given, the lexicographically next layer is assumed to be the destination. That is, in layer 1001 two convoluttions are applied, where the result of the first is forwarded to layer '1002', the second one to layer '1004'.

And here, a more complex example of a Unet with arbitrary depth, which can used in the above example of the PatchWorkModel:

#!python

def createBlock(depth=4,outK=1):
  theLayers = {}
  for z in range(depth):
    fdim = 5+z*5
    id_d = str(1000 + z+1)
    id_u = str(2000 + depth-z+1)
    theLayers[id_d+"conv"] = [{'f': conv(fdim) } , {'f': conv(fdim), 'dest':id_u+"relu" }  ]
    theLayers[id_d+"relu"] = BNrelu() + [layers.MaxPooling3D(pool_size=(2,2,2)) ]
    theLayers[id_u+"conv"] =  [conv_up(fdim)]
    theLayers[id_u+"relu"] = BNrelu()
  theLayers["3000"] =  [layers.Dropout(rate=0.5), conv(outK)]
  return CNNblock(theLayers)

Wiki

PatchWork / Home