Merge pull request #21 from jpdefrutos/master

.github/workflows/build.yml

@@ -2,9 +2,11 @@ name: Build
 
 on:
   push:
-    branches: [ master ]
+    branches:
+      - "*"
   pull_request:
-    branches: [ master ]
+    branches:
+      - "*"
 
 jobs:
   build:
@@ -72,10 +74,14 @@ jobs:
           gdown https://drive.google.com/uc?id=1shjSrFjS4PHE5sTku30PZTLPZpGu24o3
           gdown https://drive.google.com/uc?id=1bNmls5o0Rxw5HvBF1IYnEzmpysYJaywN
 
-      - name: Test inference
+      - name: Test inference .nii
         run: |
           livermask --input samples --output samples --verbose --vessels
 
+      - name: Test inference .nii.gz
+        run: |
+          livermask --input samples --output samples --verbose --vessels --extension .nii.gz
+
   conda:
     needs: build
     runs-on: ${{ matrix.os }}

README.md

@@ -53,6 +53,7 @@ livermask --input path-to-input --output path-to-output
 | `--cpu`  | to disable the GPU (force computations on CPU only) |
 | `--verbose`  | to enable verbose |
 | `--vessels` | to segment vessels |
+| `--extension` | of the output file (default: .nii) |
 
 ### Using code directly:
 If you wish to use the code directly (not as a CLI and without installing), you can run this command:
@@ -61,7 +62,7 @@ python -m livermask.livermask --input path-to-input --output path-to-output
 ```
 
 ## DICOM/NIfTI format
-Pipeline assumes input is in the NIfTI format, and output a binary volume in the same format (.nii).
+Pipeline assumes input is in the NIfTI format, and output a binary volume in the same format (.nii or .nii.gz).
 DICOM can be converted to NIfTI using the CLI [dcm2niix](https://github.com/rordenlab/dcm2niix), as such:
 ```
 dcm2niix -s y -m y -d 1 "path_to_CT_folder" "output_name"

livermask/livermask.py

@@ -54,15 +54,15 @@ def func(path, output, cpu, verbose, vessels, extension):
 
     for curr in tqdm(paths, "CT:"):
         # check if current file is a nifti file, if not, skip
-        if not curr.endswith(".nii") or not curr.endswith(".nii.gz"):
-            continue
-
-        # perform liver parenchyma segmentation, launch it in separate process to properly clear memory
-        pred = liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name, extension)
-
-        if vessels:
-            # perform liver vessel segmentation
-            vessel_segmenter(curr, output, cpu, verbose, multiple_flag, pred, name_vessel, extension)
+        if curr.endswith(".nii") or curr.endswith(".nii.gz"):
+            # perform liver parenchyma segmentation, launch it in separate process to properly clear memory
+            pred = liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name, extension)
+
+            if vessels:
+                # perform liver vessel segmentation
+                vessel_segmenter(curr, output, cpu, verbose, multiple_flag, pred, name_vessel, extension)
+        else:
+            log.info("Unsuported file: " + curr)
 
 
 def main():

livermask/utils/process.py

@@ -132,7 +132,7 @@ def liver_segmenter(params):
         raise "Caught KeyboardInterrupt, terminating worker"
 
 
-def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name_vessel):
+def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name_vessel, extension):
     # check if cupy is available, if not, set cpu=True
     try:
         import cupy
@@ -226,6 +226,6 @@ def vessel_segmenter(curr, output, cpu, verbose, multiple_flag, liver_mask, name
     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")
+        nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-vessels" + extension)
     else:
-        nib.save(resampled_lab, output + "-vessels.nii")
+        nib.save(resampled_lab, output + "-vessels" + extension)