Added option to segment the portal and hepatic veins (one class). Liver segmentation is now run in a separate Process, as TF is unable to clear the GPU manually - only when a session is killed, which is only when a Process is finished/killed.

.gitignore

@@ -4,3 +4,4 @@ dist/
 livermask.egg-info/ 
 *.h5
 *.nii
+*__pycache__/

livermask/configs/base.yml

@@ -0,0 +1,22 @@
+batchsize: 10
+iteration: 55000
+iteration_decay_start: 50000
+seed: 0
+display_interval: 100
+snapshot_interval: 5000
+evaluation_interval: 5000
+
+patch:
+  patchside: 64
+
+unet:
+  fn: model.py
+  number_of_label: 2
+
+updater:
+  fn: updater.py
+
+adam:
+  alpha: 0.0001
+  beta1: 0.9
+  beta2: 0.999

livermask/livermask.py

@@ -1,11 +1,10 @@
-import numpy as np 
+import numpy as np
 import os, sys
 from tqdm import tqdm 
 import nibabel as nib
 from nibabel.processing import resample_to_output, resample_from_to
 from scipy.ndimage import zoom
 from tensorflow.python.keras.models import load_model
-import gdown
 from skimage.morphology import remove_small_holes, binary_dilation, binary_erosion, ball
 from skimage.measure import label, regionprops
 import warnings
@@ -13,51 +12,35 @@ import argparse
 import pkg_resources
 import tensorflow as tf
 import logging as log
+import chainer
+import math
+from unet3d import UNet3D
+import yaml
+from tensorflow.keras import backend as K
+from numba import cuda
+from utils.process import liver_segmenter_wrapper, vessel_segmenter, intensity_normalization
+from utils.utils import verboseHandler
+import logging as log
+from utils.utils import get_model, get_vessel_model
 
 
 os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'  # due to this: https://github.com/tensorflow/tensorflow/issues/35029
 warnings.filterwarnings('ignore', '.*output shape of zoom.*')  # mute some warnings
 
 
-def intensity_normalization(volume, intensity_clipping_range):
-    result = np.copy(volume)
-
-    result[volume < intensity_clipping_range[0]] = intensity_clipping_range[0]
-    result[volume > intensity_clipping_range[1]] = intensity_clipping_range[1]
-
-    min_val = np.amin(result)
-    max_val = np.amax(result)
-    if (max_val - min_val) != 0:
-        result = (result - min_val) / (max_val - min_val)
-    return result
-
-def post_process(pred):
-    return pred
-
-def get_model(output):
-    url = "https://drive.google.com/uc?id=12or5Q79at2BtLgQ7IaglNGPFGRlEgEHc"
-    md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
-    gdown.cached_download(url, output, md5=md5) #, postprocess=gdown.extractall)
-
-def verboseHandler(verbose):
-    if verbose:
-        log.basicConfig(format="%(levelname)s: %(message)s", level=log.DEBUG)
-        log.info("Verbose output.")
-    else:
-        log.basicConfig(format="%(levelname)s: %(message)s")
-
-def func(path, output, cpu, verbose):
+def func(path, output, cpu, verbose, vessels):
     # enable verbose or not
-    verboseHandler(verbose)
+    log = verboseHandler(verbose)
 
     cwd = "/".join(os.path.realpath(__file__).replace("\\", "/").split("/")[:-1]) + "/"
     name = cwd + "model.h5"
+    name_vessel = cwd + "model-hepatic_vessel.npz"
 
-    # get model
+    # get models
     get_model(name)
 
-    # load model
-    model = load_model(name, compile=False)
+    if vessels:
+        get_vessel_model(name_vessel)
 
     if not os.path.isdir(path):
         paths = [path]
@@ -73,76 +56,13 @@ def func(path, output, cpu, verbose):
         if not curr.endswith(".nii"):
             continue
 
-        log.info("preprocessing...")
-        nib_volume = nib.load(curr)
-        new_spacing = [1., 1., 1.]
-        resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
-        data = resampled_volume.get_data().astype('float32')
-
-        curr_shape = data.shape
-
-        # resize to get (512, 512) output images
-        img_size = 512
-        data = zoom(data, [img_size / data.shape[0], img_size / data.shape[1], 1.0], order=1)
-
-        # intensity normalization
-        intensity_clipping_range = [-150, 250] # HU clipping limits (Pravdaray's configs)
-        data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
-
-        # fix orientation
-        data = np.rot90(data, k=1, axes=(0, 1))
-        data = np.flip(data, axis=0)
-
-        log.info("predicting...")
-        # predict on data
-        pred = np.zeros_like(data).astype(np.float32)
-        for i in tqdm(range(data.shape[-1]), "pred: ", disable=not verbose):
-            pred[..., i] = model.predict(np.expand_dims(np.expand_dims(np.expand_dims(data[..., i], axis=0), axis=-1), axis=0))[0, ..., 1]
-        del data 
-
-        # threshold
-        pred = (pred >= 0.4).astype(int)
-
-        # fix orientation back
-        pred = np.flip(pred, axis=0)
-        pred = np.rot90(pred, k=-1, axes=(0, 1))
-
-        log.info("resize back...")
-        # resize back from 512x512
-        pred = zoom(pred, [curr_shape[0] / img_size, curr_shape[1] / img_size, 1.0], order=1)
-        pred = (pred >= 0.5).astype(bool)
-
-        log.info("morphological post-processing...")
-        # morpological post-processing
-        # 1) first erode
-        pred = binary_erosion(pred, ball(3)).astype(np.int32)
-
-        # 2) keep only largest connected component
-        labels = label(pred)
-        regions = regionprops(labels)
-        area_sizes = []
-        for region in regions:
-            area_sizes.append([region.label, region.area])
-        area_sizes = np.array(area_sizes)
-        tmp = np.zeros_like(pred)
-        tmp[labels == area_sizes[np.argmax(area_sizes[:, 1]), 0]] = 1
-        pred = tmp.copy()
-        del tmp, labels, regions, area_sizes
-
-        # 3) dilate
-        pred = binary_dilation(pred.astype(bool), ball(3))
-
-        # 4) remove small holes
-        pred = remove_small_holes(pred.astype(bool), area_threshold=0.001 * np.prod(pred.shape))
-
-        log.info("saving...")
-        pred = pred.astype(np.uint8)
-        img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
-        resampled_lab = resample_from_to(img, nib_volume, order=0)
-        if multiple_flag:
-            nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-livermask" + ".nii")
-        else:
-            nib.save(resampled_lab, output + ".nii")
+        # perform liver parenchyma segmentation, launch it in separate process to properly clear memory
+        pred = liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name)
+
+        if vessels:
+            # perform liver vessel segmentation
+            vessel_segmenter(curr, output, cpu, verbose, multiple_flag, pred, name_vessel)
+
 
 def main():
     parser = argparse.ArgumentParser()
@@ -154,6 +74,8 @@ def main():
                     help="force using the CPU even if a GPU is available.")
     parser.add_argument('--verbose', action='store_true',
                     help="enable verbose.")
+    parser.add_argument('--vessels', action='store_true',
+                    help="segment vessels.")
     ret = parser.parse_args(sys.argv[1:]); print(ret)
 
     if ret.cpu:

livermask/unet3d.py

@@ -0,0 +1,93 @@
+#coding:utf-8
+'''
+* @auther mygw
+* @date 2018-6-15
+'''
+
+import chainer
+import chainer.functions as F
+import chainer.links as L
+
+class UNet3D(chainer.Chain):
+
+    def __init__(self, num_of_label):
+        w = chainer.initializers.HeNormal()
+        super(UNet3D, self).__init__()
+        with self.init_scope():
+            # encoder pass
+            self.ce0 = L.ConvolutionND(ndim=3, in_channels=1, out_channels=16, ksize=3, pad=1,initialW=w)
+            self.bne0 = L.BatchNormalization(16)
+            self.ce1 = L.ConvolutionND(ndim=3, in_channels=16, out_channels=32, ksize=3, pad=1,initialW=w)
+            self.bne1 = L.BatchNormalization(32)
+
+            self.ce2 = L.ConvolutionND(ndim=3, in_channels=32, out_channels=32, ksize=3, pad=1, initialW=w)
+            self.bne2 = L.BatchNormalization(32)
+            self.ce3 = L.ConvolutionND(ndim=3, in_channels=32, out_channels=64, ksize=3, pad=1, initialW=w)
+            self.bne3 = L.BatchNormalization(64)
+
+            self.ce4 = L.ConvolutionND(ndim=3, in_channels=64, out_channels=64, ksize=3, pad=1, initialW=w)
+            self.bne4 = L.BatchNormalization(64)
+
+            # decoder pass
+            self.cd4 = L.ConvolutionND(ndim=3, in_channels=64, out_channels=128, ksize=3, pad=1, initialW=w)
+            self.bnd4 = L.BatchNormalization(128)
+            self.deconv2 = L.DeconvolutionND(ndim=3, in_channels=128, out_channels=128, ksize=2, stride=2, initialW=w, nobias=True)
+
+            self.cd3 = L.ConvolutionND(ndim=3, in_channels=64+128, out_channels=64, ksize=3, pad=1, initialW=w)
+            self.bnd3 = L.BatchNormalization(64)
+            self.cd2 = L.ConvolutionND(ndim=3, in_channels=64, out_channels=64, ksize=3, pad=1, initialW=w)
+            self.bnd2 = L.BatchNormalization(64)
+            self.deconv1 = L.DeconvolutionND(ndim=3, in_channels=64, out_channels=64, ksize=2, stride=2, initialW=w,nobias=True)
+
+            self.cd1 = L.ConvolutionND(ndim=3, in_channels=32+64, out_channels=32, ksize=3, pad=1, initialW=w)
+            self.bnd1 = L.BatchNormalization(32)
+            self.cd0 = L.ConvolutionND(ndim=3, in_channels=32, out_channels=32, ksize=3, pad=1, initialW=w)
+            self.bnd0 = L.BatchNormalization(32)
+            self.lcl = L.ConvolutionND(ndim=3, in_channels=32, out_channels=num_of_label, ksize=1, pad=0, initialW=w)
+
+    def __call__(self, x):
+
+        # encoder pass
+        e0 = F.relu(self.bne0(self.ce0(x)))
+        e1 = F.relu(self.bne1(self.ce1(e0)))
+        del e0
+        e2 = F.relu(self.bne2(self.ce2(F.max_pooling_nd(e1, ksize=2, stride=2))))
+        e3 = F.relu(self.bne3(self.ce3(e2)))
+        del e2
+        e4 = F.relu(self.bne4(self.ce4(F.max_pooling_nd(e3, ksize=2, stride=2))))
+
+        # decoder pass
+        d4 = F.relu(self.bnd4(self.cd4(e4)))
+        del e4
+        d3 = F.relu(self.bnd3(self.cd3(F.concat([self.deconv2(d4), e3]))))
+        del d4, e3
+        d2 = F.relu(self.bnd2(self.cd2(d3)))
+        del d3
+        d1 = F.relu(self.bnd1(self.cd1(F.concat([self.deconv1(d2), e1]))))
+        del d2, e1
+        d0 = F.relu(self.bnd0(self.cd0(d1)))
+        del d1
+        lcl = F.softmax(self.lcl(d0), axis=1)
+
+        return lcl #(batchsize, ch, z, y, x)
+
+
+    def cropping(self, input, ref):
+        '''
+        * @param input encoder feature map
+        * @param ref decoder feature map
+        '''
+        edgez = (input.shape[2] - ref.shape[2])/2
+        edgey = (input.shape[3] - ref.shape[3])/2
+        edgex = (input.shape[4] - ref.shape[4])/2
+        edgez = int(edgex)
+        edgey = int(edgey)
+        edgex = int(edgez)
+
+        X = F.split_axis(input,(edgex,int(input.shape[4]-edgex)),axis=4)
+        X = X[1]
+        X = F.split_axis(X,(edgey,int(X.shape[3]-edgey)),axis=3)
+        X = X[1]
+        X = F.split_axis(X,(edgez,int (X.shape[2]-edgez)),axis=2)
+        X = X[1]
+        return X

livermask/utils/process.py

@@ -0,0 +1,210 @@
+import numpy as np
+import os, sys
+from tqdm import tqdm 
+import nibabel as nib
+from nibabel.processing import resample_to_output, resample_from_to
+from scipy.ndimage import zoom
+from tensorflow.python.keras.models import load_model
+import gdown
+from skimage.morphology import remove_small_holes, binary_dilation, binary_erosion, ball
+from skimage.measure import label, regionprops
+import warnings
+import argparse
+import pkg_resources
+import tensorflow as tf
+import logging as log
+import chainer
+import utils.yaml_utils  as yaml_utils
+import math
+from unet3d import UNet3D
+import yaml
+from tensorflow.keras import backend as K
+from numba import cuda
+from utils.utils import load_vessel_model
+import multiprocessing as mp
+
+
+def intensity_normalization(volume, intensity_clipping_range):
+    result = np.copy(volume)
+
+    result[volume < intensity_clipping_range[0]] = intensity_clipping_range[0]
+    result[volume > intensity_clipping_range[1]] = intensity_clipping_range[1]
+
+    min_val = np.amin(result)
+    max_val = np.amax(result)
+    if (max_val - min_val) != 0:
+        result = (result - min_val) / (max_val - min_val)
+    return result
+
+
+def liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name):
+    # run inference in a different process
+    p = mp.Pool(processes=1, maxtasksperchild=1)  # , initializer=initializer)
+    result = p.map_async(liver_segmenter, ((curr, output, cpu, verbose, multiple_flag, name), ))
+    return result.get()[0]
+
+
+def liver_segmenter(params):
+    try:
+        curr, output, cpu, verbose, multiple_flag, name = params
+
+        # load model
+        model = load_model(name, compile=False)
+
+        log.info("preprocessing...")
+        nib_volume = nib.load(curr)
+        new_spacing = [1., 1., 1.]
+        resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
+        data = resampled_volume.get_data().astype('float32')
+
+        curr_shape = data.shape
+
+        # resize to get (512, 512) output images
+        img_size = 512
+        data = zoom(data, [img_size / data.shape[0], img_size / data.shape[1], 1.0], order=1)
+
+        # intensity normalization
+        intensity_clipping_range = [-150, 250] # HU clipping limits (Pravdaray's configs)
+        data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
+
+        # fix orientation
+        data = np.rot90(data, k=1, axes=(0, 1))
+        data = np.flip(data, axis=0)
+
+        log.info("predicting...")
+        # predict on data
+        pred = np.zeros_like(data).astype(np.float32)
+        for i in tqdm(range(data.shape[-1]), "pred: ", disable=not verbose):
+            pred[..., i] = model.predict(np.expand_dims(np.expand_dims(np.expand_dims(data[..., i], axis=0), axis=-1), axis=0))[0, ..., 1]
+        del data 
+
+        # threshold
+        pred = (pred >= 0.4).astype(int)
+
+        # fix orientation back
+        pred = np.flip(pred, axis=0)
+        pred = np.rot90(pred, k=-1, axes=(0, 1))
+
+        log.info("resize back...")
+        # resize back from 512x512
+        pred = zoom(pred, [curr_shape[0] / img_size, curr_shape[1] / img_size, 1.0], order=1)
+        pred = (pred >= 0.5).astype(bool)
+
+        log.info("morphological post-processing...")
+        # morpological post-processing
+        # 1) first erode
+        pred = binary_erosion(pred, ball(3)).astype(np.int32)
+
+        # 2) keep only largest connected component
+        labels = label(pred)
+        regions = regionprops(labels)
+        area_sizes = []
+        for region in regions:
+            area_sizes.append([region.label, region.area])
+        area_sizes = np.array(area_sizes)
+        tmp = np.zeros_like(pred)
+        tmp[labels == area_sizes[np.argmax(area_sizes[:, 1]), 0]] = 1
+        pred = tmp.copy()
+        del tmp, labels, regions, area_sizes
+
+        # 3) dilate
+        pred = binary_dilation(pred.astype(bool), ball(3))
+
+        # 4) remove small holes
+        pred = remove_small_holes(pred.astype(bool), area_threshold=0.001 * np.prod(pred.shape))
+
+        log.info("saving...")
+        pred = pred.astype(np.uint8)
+        img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
+        resampled_lab = resample_from_to(img, nib_volume, order=0)
+        if multiple_flag:
+            nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-livermask.nii")
+        else:
+            nib.save(resampled_lab, output  + "-livermask.nii")
+
+        return pred
+    except KeyboardInterrupt:
+        raise "Caught KeyboardInterrupt, terminating worker"
+
+
+def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name_vessel):
+
+    # load model
+    unet, xp = load_vessel_model(name_vessel, cpu)
+
+    # read config
+    config = yaml_utils.Config(yaml.safe_load(open("./configs/base.yml")))
+
+    # read data
+    nib_volume = nib.load(curr)
+    new_spacing = [1., 1., 1.]
+    resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
+    #resampled_volume = nib_volume
+    org = resampled_volume.get_data().astype('float32')
+
+    # HU clipping
+    intensity_clipping_range = [80, 220]
+    org[org < intensity_clipping_range[0]] = intensity_clipping_range[0]
+    org[org > intensity_clipping_range[1]] = intensity_clipping_range[1]
+
+    # Calculate maximum of number of patch at each side
+    ze,ye,xe = org.shape
+    xm = int(math.ceil((float(xe)/float(config.patch['patchside']))))
+    ym = int(math.ceil((float(ye)/float(config.patch['patchside']))))
+    zm = int(math.ceil((float(ze)/float(config.patch['patchside']))))
+
+    margin = ((0, config.patch['patchside']),
+              (0, config.patch['patchside']),
+              (0, config.patch['patchside']))
+    org = np.pad(org, margin, 'edge')
+    org = chainer.Variable(xp.array(org[np.newaxis, np.newaxis, :], dtype=xp.float32))
+
+    # init prediction array
+    prediction_map = np.zeros((ze + config.patch['patchside'], ye + config.patch['patchside'], xe + config.patch['patchside']))
+    probability_map = np.zeros((config.unet['number_of_label'], ze+config.patch['patchside'], ye + config.patch['patchside'], xe + config.patch['patchside']))
+
+    # Patch loop
+    for s in tqdm(range(xm * ym * zm), ascii=True, desc='Patch loop'):
+        xi = int(s % xm) * config.patch['patchside']
+        yi = int((s % (ym * xm)) / xm) * config.patch['patchside']
+        zi = int(s / (ym * xm)) * config.patch['patchside']
+
+        # check if current region contains any liver mask, if not, skip
+        parenchyma_patch = liver_mask[zi:zi + config.patch['patchside'], yi:yi + config.patch['patchside'], xi:xi + config.patch['patchside']]
+        #if np.count_nonzero(parenchyma_patch) == 0:
+        if np.mean(parenchyma_patch) < 0.25:
+            continue
+
+        # Extract patch from original image
+        patch = org[:, :, zi:zi + config.patch['patchside'], yi:yi + config.patch['patchside'], xi:xi + config.patch['patchside']]
+        with chainer.using_config('train', False), chainer.using_config('enable_backprop', False):
+            probability_patch = unet(patch)
+
+        # Generate probability map
+        probability_patch = probability_patch.data
+        #if args.gpu >= 0:
+        if not cpu:
+            probability_patch = chainer.cuda.to_cpu(probability_patch)
+        for ch in range(probability_patch.shape[1]):
+            probability_map[ch, zi:zi + config.patch['patchside'],yi:yi + config.patch['patchside'], xi:xi + config.patch['patchside']] = probability_patch[0, ch, :, :, :]
+
+        prediction_patch = np.argmax(probability_patch, axis=1)
+
+        prediction_map[zi:zi + config.patch['patchside'], yi:yi + config.patch['patchside'], xi:xi + config.patch['patchside']] = prediction_patch[0, :, :, :]
+
+    print('Save image')
+    probability_map = probability_map[:, :ze, :ye, :xe]
+    prediction_map = prediction_map[:ze, :ye, :xe]
+
+    # post-process prediction
+    #prediction_map = prediction_map + liver_mask
+    #prediction_map[prediction_map > 0] = 1
+
+    log.info("saving...")
+    pred = prediction_map.astype(np.uint8)
+    img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
+    resampled_lab = resample_from_to(img, nib_volume, order=0)
+    if multiple_flag:
+        nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-vessels.nii")
+    else:
+        nib.save(resampled_lab, output + "-vessels.nii")

livermask/utils/utils.py

@@ -0,0 +1,38 @@
+import gdown
+import logging as log
+import chainer
+from unet3d import UNet3D
+
+
+def get_model(output):
+    url = "https://drive.google.com/uc?id=12or5Q79at2BtLgQ7IaglNGPFGRlEgEHc"
+    md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
+    gdown.cached_download(url, output, md5=md5) #, postprocess=gdown.extractall)
+
+
+def get_vessel_model(output):
+    url = "https://drive.google.com/uc?id=1-8VNoRmIeiF1uIuWBqmZXz_6dIQFSAxN"
+    #md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
+    gdown.cached_download(url, output) #, md5=md5)
+
+
+def load_vessel_model(path, cpu):
+    unet = UNet3D(num_of_label=2)
+    chainer.serializers.load_npz(path, unet)
+
+    if not cpu:
+        chainer.cuda.get_device_from_id(0).use()
+        unet.to_gpu()
+
+    xp = unet.xp
+
+    return unet, xp
+
+
+def verboseHandler(verbose):
+    if verbose:
+        log.basicConfig(format="%(levelname)s: %(message)s", level=log.DEBUG)
+        log.info("Verbose output.")
+    else:
+        log.basicConfig(format="%(levelname)s: %(message)s")
+    return log

livermask/utils/yaml_utils.py

@@ -0,0 +1,28 @@
+#coding:utf-8
+import shutil
+import sys, os, time
+import yaml
+from utils import yaml_utils
+
+
+#sys.path.append(os.path.dirname(__file__))
+
+# Copy from tgans repo.
+class Config(object):
+    '''
+    'https://github.com/pfnet-research/sngan_projection/blob/master/source/yaml_utils.py'
+    '''
+    def __init__(self, config_dict):
+        self.config = config_dict
+
+    def __getattr__(self, key):
+        if key in self.config:
+            return self.config[key]
+        else:
+            raise AttributeError(key)
+
+    def __getitem__(self, key):
+        return self.config[key]
+
+    def __repr__(self):
+        return yaml.dump(self.config, default_flow_style=False)