refactored to PEP8

andreped · Oct 27, 2022 · 30d661c · 30d661c
1 parent cc8381b
commit 30d661c
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 61 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -88,10 +88,6 @@ jobs:
 
       - name: Install wheel
         run: |
-          python -V
-          python3 -V
-          pip -V
-          pip3 -V
           pip3 install --find-links=${{github.workspace}} livermask
 
       - name: Download test sample

diff --git a/livermask/livermask.py b/livermask/livermask.py
@@ -1,5 +1,6 @@
 import numpy as np
-import os, sys
+import os
+import sys
 from tqdm import tqdm 
 import nibabel as nib
 from nibabel.processing import resample_to_output, resample_from_to

diff --git a/livermask/utils/process.py b/livermask/utils/process.py
@@ -1,6 +1,6 @@
 import numpy as np
 import os, sys
-from tqdm import tqdm 
+from tqdm import tqdm
 import nibabel as nib
 from nibabel.processing import resample_to_output, resample_from_to
 from scipy.ndimage import zoom
@@ -41,7 +41,7 @@ def liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name):
     # run inference in a different process
     mp.set_start_method('spawn', force=True)
     with mp.Pool(processes=1, maxtasksperchild=1) as p:  # , initializer=initializer)
-        result = p.map_async(liver_segmenter, ((curr, output, cpu, verbose, multiple_flag, name), ))
+        result = p.map_async(liver_segmenter, ((curr, output, cpu, verbose, multiple_flag, name),))
         log.info("getting result from process...")
         ret = result.get()[0]
     return ret
@@ -67,7 +67,7 @@ def liver_segmenter(params):
         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)
+        intensity_clipping_range = [-150, 250]  # HU clipping limits (Pravdaray's configs)
         data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
 
         # fix orientation
@@ -78,8 +78,10 @@ def liver_segmenter(params):
         # 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 
+            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)
@@ -123,15 +125,14 @@ def liver_segmenter(params):
         if multiple_flag:
             nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-livermask.nii")
         else:
-            nib.save(resampled_lab, output  + "-livermask.nii")
+            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):
-
     # check if cupy is available, if not, set cpu=True
     try:
         import cupy
@@ -150,7 +151,7 @@ def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name
     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
+    # resampled_volume = nib_volume
     org = resampled_volume.get_data().astype('float32')
 
     # HU clipping
@@ -159,10 +160,10 @@ def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name
     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']))))
+    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']),
@@ -171,8 +172,10 @@ def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name
     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']))
+    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']))
 
     log.info("predicting...")
     # Patch loop
@@ -182,34 +185,38 @@ def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name
         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:
+        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']]
+        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 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, :, :, :]
+            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, :, :, :]
+        prediction_map[zi:zi + config.patch['patchside'], yi:yi + config.patch['patchside'],
+                       xi:xi + config.patch['patchside']] = prediction_patch[0, :, :, :]
 
-    probability_map = probability_map[:, :ze, :ye, :xe]
+    # 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
+    # prediction_map = prediction_map + liver_mask
+    # prediction_map[prediction_map > 0] = 1
 
     # filter segmented vessels outside the predicted liver parenchyma
     pred = prediction_map.astype(np.uint8)

diff --git a/livermask/utils/unet3d.py b/livermask/utils/unet3d.py
@@ -1,23 +1,17 @@
-#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):
 
+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.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.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)
@@ -31,22 +25,23 @@ def __init__(self, num_of_label):
             # 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.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.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.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.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)))
@@ -69,25 +64,24 @@ def __call__(self, x):
         del d1
         lcl = F.softmax(self.lcl(d0), axis=1)
 
-        return lcl #(batchsize, ch, z, y, x)
-
+        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 = (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 = 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 = 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 = F.split_axis(X, (edgez, int(X.shape[2] - edgez)), axis=2)
         X = X[1]
         return X
diff --git a/livermask/utils/utils.py b/livermask/utils/utils.py
@@ -6,14 +6,12 @@
 
 def get_model(output):
     url = "https://drive.google.com/uc?id=12or5Q79at2BtLgQ7IaglNGPFGRlEgEHc"
-    md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
-    gdown.cached_download(url, output) #, md5=md5) #, postprocess=gdown.extractall)
+    gdown.cached_download(url, output)
 
 
 def get_vessel_model(output):
     url = "https://drive.google.com/uc?id=1-8VNoRmIeiF1uIuWBqmZXz_6dIQFSAxN"
-    #md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
-    gdown.cached_download(url, output) #, md5=md5)
+    gdown.cached_download(url, output)
 
 
 def load_vessel_model(path, cpu):
@@ -22,7 +20,6 @@ def load_vessel_model(path, cpu):
     if not cpu:
         chainer.cuda.get_device_from_id(0).use()
         unet.to_gpu()
-
     xp = unet.xp
     return unet, xp
 

diff --git a/livermask/utils/yaml_utils.py b/livermask/utils/yaml_utils.py
@@ -1,13 +1,8 @@
-import shutil
-import sys, os, time
 import yaml
 
 
-# Copy from tgans repo.
+# https://github.com/pfnet-research/sngan_projection/blob/master/source/yaml_utils.py
 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