Add brain extraction step

README.md

@@ -26,6 +26,24 @@ or locally on the command line. To use it locally, move input NIFTI files to a d
 apptainer run docker://ghcr.io/fnndsc/pl-emerald:latest emerald input/ output/
 ```
 
+To create masks next to the original file, with the names `*_mask.nii`:
+
+```shell
+apptainer run docker://ghcr.io/fnndsc/pl-emerald:latest emerald --mask-suffix _mask.nii input/ input/
+```
+
+To extract brains without keeping the mask file:
+
+```shell
+apptainer run docker://ghcr.io/fnndsc/pl-emerald:latest emerald --mask-suffix '' --outputs '0:.nii' input/ output/
+```
+
+To create output masks, extracted brains, and masks overlayed on the original with dimmed background (for convenient visualization):
+
+```shell
+apptainer run docker://ghcr.io/fnndsc/pl-emerald:latest emerald --mask-suffix '_mask.nii' --outputs '0.0:_brain.nii,0.2:_overlay02.nii' input/ output/
+```
+
 ## Limitations
 
 - Unet can currently only work with 256x256 images

emerald/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.1'
+__version__ = '0.2.0'
 
 DISPLAY_TITLE = r"""
        _                                     _     _ 

emerald/__main__.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python
-
+import sys
 from pathlib import Path
 from argparse import ArgumentParser, Namespace, ArgumentDefaultsHelpFormatter
+from typing import Optional, List, Tuple
 
 from chris_plugin import chris_plugin, PathMapper, curry_name_mapper
 
@@ -19,8 +20,10 @@
                         formatter_class=ArgumentDefaultsHelpFormatter)
 parser.add_argument('-p', '--pattern', type=str, default='**/*.nii',
                     help='Input files pattern')
-parser.add_argument('-s', '--output-suffix', type=str, default='_mask.nii',
-                    help='Output file suffix')
+parser.add_argument('-m', '--mask-suffix', type=str, default='_mask.nii',
+                    help='Mask output file suffix. Provide "" to not save mask.')
+parser.add_argument('-o', '--outputs', type=str, default='',
+                    help='Background intensity multiplier and output suffix.')
 parser.add_argument('--no-post-processing', dest='post_processing', action='store_false',
                     help='Predicted mask should not be post processed (morphological closing and defragmentation)')
 parser.add_argument('--dilation-footprint', default='disk(2)', type=str,
@@ -40,11 +43,38 @@ def main(options: Namespace, inputdir: Path, outputdir: Path):
 
     model = Unet()
     footprint = eval(options.dilation_footprint)
+    outputs = parse_outputs(options.outputs)
 
-    mapper = PathMapper.file_mapper(inputdir, outputdir,
-                                    glob=options.pattern, name_mapper=curry_name_mapper('{}_mask.nii'))
+    mapper = PathMapper.file_mapper(inputdir, outputdir, glob=options.pattern)
     for input_file, output_file in mapper:
-        emerald(model, input_file, output_file, options.post_processing, footprint)
+        mask_path = change_suffix(output_file, options.mask_suffix)
+        brain_path = [(n, change_suffix(output_file, s)) for n, s in outputs]
+        emerald(model, input_file, mask_path, brain_path, options.post_processing, footprint)
+
+
+def change_suffix(path: Path, suffix: Optional[str]) -> Optional[Path]:
+    if not suffix:
+        return None
+    if '.' not in path.name:
+        return path.with_name(path.name + suffix)
+    name_part, _old_suffix = path.name.rsplit('.', maxsplit=1)
+    return path.with_name(name_part + suffix)
+
+
+def parse_outputs(val: str) -> List[Tuple[float, str]]:
+    val = val.strip()
+    if not val:
+        return []
+    try:
+        return [parse_pair(p) for p in val.split(',')]
+    except ValueError as e:
+        print(e)
+        sys.exit(1)
+
+
+def parse_pair(val: str) -> Tuple[float, str]:
+    num, suffix = val.split(':', maxsplit=1)
+    return float(num), suffix
 
 
 if __name__ == '__main__':

emerald/emerald.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Optional
+from typing import Optional, List, Tuple
 
 import cv2
 import numpy as np
@@ -97,7 +97,8 @@ def __postProcessing(mask, no_dilation, footprint):
     return pred_mask
 
 
-def emerald(model: Unet, input_path: str, output_path: Path, post_processing: bool, footprint: Optional[npt.NDArray]):
+def emerald(model: Unet, input_path: str, mask_path: Optional[Path], brain_paths: List[Tuple[float, Path]],
+            post_processing: bool, footprint: Optional[npt.NDArray]):
     img_path = str(input_path)
 
     img, hdr = getImageData(img_path)
@@ -114,6 +115,7 @@ def emerald(model: Unet, input_path: str, output_path: Path, post_processing: bo
         res = __postProcessing(res, no_dilation=(footprint is not None), footprint=footprint)
 
     if resizeNeeded:
+        # jennings to sofia: why np.float32 instead of uint8?
         res = __resizeData(res.astype(np.float32), target = original_shape)
 
     #remove extra dimension
@@ -123,4 +125,17 @@ def emerald(model: Unet, input_path: str, output_path: Path, post_processing: bo
     res = np.moveaxis(res, 0, -1)
 
     #save result
-    save(res, str(output_path), hdr)
+    if mask_path:
+        save(res, str(mask_path), hdr)
+
+    if brain_paths:
+        # for whatever reason, img.shape=(38, 256, 256, 1).
+        if len(img.shape) == 4 and img.shape[3] == 1:
+            img = np.squeeze(img)
+            img = np.moveaxis(img, 0, -1)
+
+    # apply res mask to img
+    for mult, brain_path in brain_paths:
+        print(f'img.shape={img.shape}, res.shape={res.shape}')
+        overlayed_data = np.clip(res, mult, 1.0) * img
+        save(overlayed_data, str(brain_path), hdr)