From fa4a47f0e88258a7d226ee4afee8974c694c9786 Mon Sep 17 00:00:00 2001 From: Jennings Zhang Date: Tue, 14 Nov 2023 03:27:06 -0500 Subject: [PATCH] Add brain extraction step --- README.md | 18 ++++++++++++++++++ emerald/__init__.py | 2 +- emerald/__main__.py | 42 ++++++++++++++++++++++++++++++++++++------ emerald/emerald.py | 21 ++++++++++++++++++--- 4 files changed, 73 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index e5d5ffc..ef17b2d 100644 --- a/README.md +++ b/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 diff --git a/emerald/__init__.py b/emerald/__init__.py index 1542922..c8cf5f3 100644 --- a/emerald/__init__.py +++ b/emerald/__init__.py @@ -1,4 +1,4 @@ -__version__ = '0.1.1' +__version__ = '0.2.0' DISPLAY_TITLE = r""" _ _ _ diff --git a/emerald/__main__.py b/emerald/__main__.py index 9a99e07..7302d4a 100644 --- a/emerald/__main__.py +++ b/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__': diff --git a/emerald/emerald.py b/emerald/emerald.py index 19fa176..4cd26cb 100644 --- a/emerald/emerald.py +++ b/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)