diff --git a/models/bamf_pet_ct_lung_tumor/config/default.yml b/models/bamf_pet_ct_lung_tumor/config/default.yml new file mode 100644 index 00000000..d4262634 --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/config/default.yml @@ -0,0 +1,48 @@ +general: + data_base_dir: /app/data + version: 1.0 + description: default configuration for Bamf NNUnet Lung and FDG-avid lesions in the lung from FDG PET/CT scans (dicom to dicom) + +execute: +- FileStructureImporter +- NiftiConverter +- Registration +- NNUnetPETCTRunner +- LungSegmentatorRunner +- LungPostProcessor +- DsegConverter +- DataOrganizer + +modules: + FileStructureImporter: + input_dir: 'input_data' + structures: + - $patientID@instance/ct@dicom:mod=ct + - $patientID/pt@dicom:mod=pt + import_id: patientID + + NiftiConverter: + in_datas: dicom:mod=pt|ct + engine: dcm2niix + allow_multi_input: true + + NNUnetPETCTRunner: + in_ct_data: nifti:mod=ct:registered=true + in_pt_data: nifti:mod=pt + nnunet_task: Task762_PET_CT_Breast + nnunet_model: 3d_fullres + roi: LIVER,KIDNEY,URINARY_BLADDER,SPLEEN,LUNG,BRAIN,HEART,SMALL_INTESTINE,LUNG+FDG_AVID_TUMOR + + LungSegmentatorRunner: + in_data: nifti:mod=ct:registered=true + use_fast_mode: True + + DsegConverter: + source_segs: nifti:mod=seg:processor=bamf + model_name: BAMF Lung and FDG Tumor Segmentation + target_dicom: dicom:mod=pt + skip_empty_slices: True + + DataOrganizer: + targets: + - dicomseg-->[i:patientID]/bamf_pet_ct_lung_tumor.seg.dcm \ No newline at end of file diff --git a/models/bamf_pet_ct_lung_tumor/dockerfiles/Dockerfile b/models/bamf_pet_ct_lung_tumor/dockerfiles/Dockerfile new file mode 100644 index 00000000..90ffd52a --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/dockerfiles/Dockerfile @@ -0,0 +1,32 @@ +FROM mhubai/base:latest + +# FIXME: set this environment variable as a shortcut to avoid nnunet crashing the build +# by pulling sklearn instead of scikit-learn +# N.B. this is a known issue: +# https://github.com/MIC-DKFZ/nnUNet/issues/1281 +# https://github.com/MIC-DKFZ/nnUNet/pull/1209 +ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True + +# Install nnunet and totalsegmentator +RUN pip3 install p_tqdm==1.4 TotalSegmentator==1.5.7 nnunet==1.6.6 --no-cache-dir + +# Clone the main branch of MHubAI/models +ARG MHUB_MODELS_REPO +RUN buildutils/import_mhub_model.sh bamf_pet_ct_lung_tumor ${MHUB_MODELS_REPO} + +# Pull nnUNet model weights into the container for Task777_CT_Nodules +ENV WEIGHTS_DIR=/root/.nnunet/nnUNet_models/nnUNet/ +RUN mkdir -p $WEIGHTS_DIR +ENV TASK_NAME=Task762_PET_CT_Breast +ENV WEIGHTS_FN=$TASK_NAME.zip +ENV WEIGHTS_URL=https://zenodo.org/record/8290055/files/$WEIGHTS_FN +RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL} --no-check-certificate +RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR} +RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN} + +# specify nnunet specific environment variables +ENV WEIGHTS_FOLDER=$WEIGHTS_DIR + +# Default run script +ENTRYPOINT ["mhub.run"] +CMD ["--config", "/app/models/bamf_pet_ct_lung_tumor/config/default.yml"] diff --git a/models/bamf_pet_ct_lung_tumor/meta.json b/models/bamf_pet_ct_lung_tumor/meta.json new file mode 100644 index 00000000..2a3c4e7a --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/meta.json @@ -0,0 +1,185 @@ +{ + "id": "", + "name": "bamf_pet_ct_lung_tumor", + "title": "FDG PET/CT Lung and Lung Tumor Annotation", + "summary": { + "description": "An nnU-Net based model to segment Lung and FDG-avid lesions in the lung from FDG PET/CT scans", + "inputs": [ + { + "label": "Input Image", + "description": "The CT scan of a patient.", + "format": "DICOM", + "modality": "CT", + "bodypartexamined": "Lung", + "slicethickness": "5mm", + "non-contrast": true, + "contrast": false + }, + { + "label": "Input Image", + "description": "The PET scan of a patient.", + "format": "DICOM", + "modality": "PT", + "bodypartexamined": "Lung", + "slicethickness": "4mm", + "non-contrast": false, + "contrast": false + } + ], + "outputs": [ + { + "label": "Segmentation", + "type": "Segmentation", + "description": "Segmentation Lung and tumor", + "classes": [ + "LUNG", + "LUNG+FDG_AVID_TUMOR" + ] + } + ], + "model": { + "architecture": "U-net", + "training": "supervised", + "cmpapproach": "3D" + }, + "data": { + "training": { + "vol_samples": 1014 + }, + "evaluation": { + "vol_samples": 77 + }, + "public": true, + "external": true + } + }, + "details": { + "name": "AIMI PET CT Lung", + "version": "1.0.0", + "devteam": "BAMF Health", + "authors": [ + "Soni, Rahul", + "McCrumb, Diana", + "Murugesan, Gowtham Krishnan", + "Van Oss, Jeff" + ], + "type": "nnU-Net (U-Net structure, optimized by data-driven heuristics)", + "date": { + "code": "17.10.2023", + "weights": "28.08.2023", + "pub": "23.10.2023" + }, + "cite": "Murugesan, Gowtham Krishnan, Diana McCrumb, Mariam Aboian, Tej Verma, Rahul Soni, Fatima Memon, and Jeff Van Oss. The AIMI Initiative: AI-Generated Annotations for Imaging Data Commons Collections. arXiv preprint arXiv:2310.14897 (2023).", + "license": { + "code": "MIT", + "weights": "CC BY-NC 4.0" + }, + "publications": [ + { + "title": "The AIMI Initiative: AI-Generated Annotations in IDC Collections", + "uri": "https://arxiv.org/abs/2310.14897" + } + ], + "github": "https://github.com/bamf-health/aimi-lung-pet-ct" + }, + "info": { + "use": { + "title": "Intended Use", + "text": "This model is intended to perform lung and tumor segmentation in FDG PET CT scans. The model has been trained and tested on scans acquired during clinical care of patients, so it might not be suited for a healthy population. The generalization capabilities of the model on a range of ages, genders, and ethnicities are unknown. For detailed information on the training set design, please refer to reference section in the training section" + }, + "analyses": { + "title": "Quantitative Analyses", + "text": "Label-wise metrics (mean (standard deviation)) between AI derived and manual corrected FDG PET/CT lungs and tumor annotations.", + "tables": [ + { + "label": "Segmentation Metric Expert", + "entries": { + "Lung DSC": "1.00 (0.00) ", + "Tumor DSC": "0.97 (0.11) ", + "Lung 95% Hausdorff (mm)": "0.10 (0.58)", + "Tumor 95% Hausdorff (mm)": "5.83 (19.42) " + } + }, + { + "label": "Segmentation Metric Tumor - Non-Expert", + "entries": { + "Lung DSC": "0.99 (0.04)", + "Tumor DSC": "0.92 (0.20)", + "Lung 95% Hausdorff (mm)": "1.97 (10.50)", + "Tumor 95% Hausdorff (mm)": "10.00 (26.34)" + } + }, + { + "label": "Detection Accuracy", + "entries": { + "Sensitivity": "0.91", + "False negative rate": "0.09", + "F1 score": "0.94" + } + } + ], + "references": [ + { + "label": "The AIMI Initiative: AI-Generated Annotations for Imaging Data Commons Collections", + "uri": "https://arxiv.org/abs/2310.14897" + } + ] + }, + "evaluation": { + "title": "Evaluation Data", + "text": "The model was used to segment cases from the IDC [1] collection of ACRIN-NSCLC-FDG-PET [2], Anti-PD-1-Lung [3], LUNG-PET-CT-Dx[4], NSCLC Radiogenomics[5], RIDER Lung PET-CT[6], TCGA-LUAD[7], TCGA-LUSC[8] . Approximately 20% of those cases were randomly selected to be reviewed and corrected by a board-certified radiologist. The model predictions, and radiologist corrections are published on zenodo [9]", + "references": [ + { + "label": "Imaging Data Collections (IDC)", + "uri": "https://datacommons.cancer.gov/repository/imaging-data-commons" + }, + { + "label": "ACRIN-NSCLC-FDG-PET", + "uri": "https://www.cancerimagingarchive.net/collection/acrin-nsclc-fdg-pet/" + }, + { + "label": "Anti-PD-1-Lung", + "uri": "https://www.cancerimagingarchive.net/collection/anti-pd-1_lung/" + }, + { + "label": "LUNG-PET-CT-Dx", + "uri": "https://www.cancerimagingarchive.net/collection/lung-pet-ct-dx/" + }, + { + "label": "NSCLC Radiogenomics", + "uri": "https://www.cancerimagingarchive.net/collection/nsclc-radiogenomics/" + }, + { + "label": "RIDER Lung PET-CT", + "uri": "https://wiki.cancerimagingarchive.net/display/Public/RIDER+Collections" + }, + { + "label": "TCGA-LUAD", + "uri": "https://www.cancerimagingarchive.net/collection/tcga-luad/" + }, + { + "label": "TCGA-LUSC", + "uri": "https://www.cancerimagingarchive.net/collection/tcga-lusc/" + }, + { + "label": "Image segmentations produced by the AIMI Annotations initiative", + "uri": "https://zenodo.org/records/10009368" + } + ] + }, + "training": { + "title": "Training Data", + "text": "The AutoPET Challenge 2023 dataset is comprised of whole-body FDG-PET/CT data from 900 patients, encompassing 1014 studies with tumor annotations. This dataset was augmented by adding labels for the bladder, kidneys, liver, stomach, spleen, lungs, and heart generated by the TotalSegmentator model. A multi-task AI model was trained using the augmented datasets", + "references": [ + { + "label": "AutoPET Challenge 2023 dataset", + "uri": "https://doi.org/10.7937/gkr0-xv29" + }, + { + "label": "Total Segmentator", + "uri": "https://doi.org/10.48550/arXiv.2208.05868" + } + ] + } + } +} \ No newline at end of file diff --git a/models/bamf_pet_ct_lung_tumor/mhub.toml b/models/bamf_pet_ct_lung_tumor/mhub.toml new file mode 100644 index 00000000..89a19666 --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/mhub.toml @@ -0,0 +1,3 @@ + +[model.deployment] +test = "https://zenodo.org/records/13880663/files/bamf_pet_ct_lung_tumor.test.zip?download=1" \ No newline at end of file diff --git a/models/bamf_pet_ct_lung_tumor/utils/LungPostProcessor.py b/models/bamf_pet_ct_lung_tumor/utils/LungPostProcessor.py new file mode 100644 index 00000000..56106551 --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/utils/LungPostProcessor.py @@ -0,0 +1,140 @@ +""" +--------------------------------------------------------- +Post processing Module on segmentation output +--------------------------------------------------------- + +------------------------------------------------- +Author: Jithendra Kumar +Email: Jithendra.kumar@bamfhealth.com +------------------------------------------------- + +""" +import os +import shutil +import SimpleITK as sitk +import numpy as np +from skimage import measure +from mhubio.core import IO +from mhubio.core import Module, Instance, InstanceData, InstanceDataCollection + + +class LungPostProcessor(Module): + + def n_connected(self, img_data): + """ + Get the largest connected component in a binary image. + + Args: + img_data (np.ndarray): image data. + + Returns: + np.ndarray: Processed image with the largest connected component. + """ + img_data_mask = np.zeros(img_data.shape) + img_data_mask[img_data >= 1] = 1 + img_filtered = np.zeros(img_data_mask.shape) + blobs_labels = measure.label(img_data_mask, background=0) + lbl, counts = np.unique(blobs_labels, return_counts=True) + lbl_dict = {} + for i, j in zip(lbl, counts): + lbl_dict[i] = j + sorted_dict = dict(sorted(lbl_dict.items(), key=lambda x: x[1], reverse=True)) + count = 0 + + for key, value in sorted_dict.items(): + if count == 1: + print(key, value) + img_filtered[blobs_labels == key] = 1 + count += 1 + + img_data[img_filtered != 1] = 0 + return img_data + + def arr_2_sitk_img(self, arr, ref): + """ + Convert numpy array to SimpleITK image. + + Args: + arr (np.ndarray): Input image data as a numpy array. + ref: Reference image for copying information. + + Returns: + sitk.Image: Converted SimpleITK image. + """ + op_img = sitk.GetImageFromArray(arr) + op_img.CopyInformation(ref) + return op_img + + def get_mets(self, left, op_data): + """ + Perform metastasis segmentation. + + Args: + left (np.ndarray): Image data for left lung. + op_data (np.ndarray): Image data for segmented regions. + + Returns: + np.ndarray: Metastasis segmentation results. + """ + op_data[left == 1] = 0 + op_primary = self.n_connected(np.copy(op_data)) + mets = np.zeros(op_primary.shape) + mets[op_data > 0] = 1 + mets[op_primary > 0] = 0 + return mets + + def get_lung_segments(self, img_path): + """ + Perform lung tissue segmentation. + + Args: + img_path (str): Path to the image for lung tissue segmentation. + + Returns: + tuple: A tuple containing lung segmentation results. + """ + img_data = sitk.GetArrayFromImage(sitk.ReadImage(img_path)) + lung_left = np.zeros(img_data.shape) + lung_left[img_data==1] = 1 + lung_right = np.zeros(img_data.shape) + lung_right[img_data==2] = 1 + return lung_left, lung_right, lung_right + lung_left + + @IO.Instance() + @IO.Input('in_ct_data', 'nifti:mod=ct:registered=true', the='input ct data') + @IO.Input('in_tumor_data', 'nifti:mod=seg:model=nnunet', the='input tumor segmentation') + @IO.Input('in_lung_seg_data', 'nifti:mod=seg:model=LungSegmentator', the='input lung segmentation') + @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=LUNG,LUNG+FDG_AVID_TUMOR', + data='in_tumor_data', + the="get the lung and tumor after post processing") + def task(self, instance: Instance, in_ct_data: InstanceData, in_tumor_data: InstanceData, + in_lung_seg_data: InstanceData, out_data: InstanceData): + """ + Perform postprocessing and writes simpleITK Image + """ + self.v("Running LungPostprocessor.") + tumor_seg_path = in_tumor_data.abspath + lung_seg_path = in_lung_seg_data.abspath + + right, left, lung = self.get_lung_segments(str(lung_seg_path)) + tumor_label = 9 + tumor_arr = sitk.GetArrayFromImage(sitk.ReadImage(tumor_seg_path)) + tumor_arr[tumor_arr != tumor_label] = 0 + + op_data = np.zeros(lung.shape) + ref = sitk.ReadImage(in_ct_data.abspath) + ct_data = sitk.GetArrayFromImage(ref) + op_data[lung == 1] = 1 + op_data[tumor_arr > 0] = 2 + th = np.min(ct_data) + op_data[ct_data == th] = 0 # removing predictions where CT not available + mets_right = self.get_mets(left, np.copy(op_data)) + mets_left = self.get_mets(right, np.copy(op_data)) + mets = np.logical_and(mets_right, mets_left).astype("int") + op_data[mets == 1] = 3 + op_data[op_data == 3] = 0 + + op_img = sitk.GetImageFromArray(op_data) + op_img.CopyInformation(ref) + + sitk.WriteImage(op_img, out_data.abspath) \ No newline at end of file diff --git a/models/bamf_pet_ct_lung_tumor/utils/LungSegmentatorRunner.py b/models/bamf_pet_ct_lung_tumor/utils/LungSegmentatorRunner.py new file mode 100644 index 00000000..3f769169 --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/utils/LungSegmentatorRunner.py @@ -0,0 +1,137 @@ +""" +------------------------------------------------- +MHub - Run Lung segmentator Module using TotalSegmentator. +------------------------------------------------- + +------------------------------------------------- +Author: Jithendra Kumar +Email: Jithendra.kumar@bamfhealth.com +------------------------------------------------- +""" + +from mhubio.core import Module, Instance, InstanceData, DataType, FileType, CT, SEG, IO, DataTypeQuery +import os, subprocess +import SimpleITK as sitk +import numpy as np +from skimage import measure +from mhubio.core import IO + +from totalsegmentator.map_to_binary import class_map + +@IO.ConfigInput('in_data', 'nifti:mod=ct', the="input data to run Lung Segmentator on") +@IO.Config('use_fast_mode', bool, True, the="flag to set to run TotalSegmentator in fast mode") +class LungSegmentatorRunner(Module): + + use_fast_mode: bool + + def mask_labels(self, labels, ts): + """ + Create a mask based on given labels. + + Args: + labels (list): List of labels to be masked. + ts (np.ndarray): Image data. + + Returns: + np.ndarray: Masked image data. + """ + lung = np.zeros(ts.shape) + for lbl in labels: + lung[ts == lbl] = 1 + return lung + + def combine_labels(self, file_paths): + """ + Create a combined segment from list of segment files. + + Args: + file_paths (list): List of segment files. + + Returns: + np.ndarray: Combined segment. + """ + images = [sitk.ReadImage(file_path) for file_path in file_paths if os.path.exists(file_path)] + result_image = sitk.GetArrayFromImage(images[0]) + # Combine segments by summing + for img in images[1:]: + img_arr = sitk.GetArrayFromImage(img) + result_image = result_image + img_arr + segment = np.zeros(result_image.shape) + segment[result_image > 0] = 1 + return segment + + def n_connected(self, img_data): + """ + Get the largest connected component in a binary image. + + Args: + img_data (np.ndarray): image data. + + Returns: + np.ndarray: Processed image with the largest connected component. + """ + img_data_mask = np.zeros(img_data.shape) + img_data_mask[img_data >= 1] = 1 + img_filtered = np.zeros(img_data_mask.shape) + blobs_labels = measure.label(img_data_mask, background=0) + lbl, counts = np.unique(blobs_labels, return_counts=True) + lbl_dict = {} + for i, j in zip(lbl, counts): + lbl_dict[i] = j + sorted_dict = dict(sorted(lbl_dict.items(), key=lambda x: x[1], reverse=True)) + count = 0 + + for key, value in sorted_dict.items(): + if count == 1: + print(key, value) + img_filtered[blobs_labels == key] = 1 + count += 1 + + img_data[img_filtered != 1] = 0 + return img_data + + + @IO.Instance() + @IO.Input('in_data', the="input whole body ct scan") + @IO.Output('out_data', 'lung_segmentations.nii.gz', 'nifti:mod=seg:model=LungSegmentator:roi=LEFT_LUNG,RIGHT_LUNG', data='in_data', the="output segmentation mask containing lung labels") + def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None: + # use total segmentator to extract lung labels + bash_command = ["TotalSegmentator"] + bash_command += ["-i", in_data.abspath] + + tmp_dir = self.config.data.requestTempDir(label="lung-segment-processor") + bash_command += ["-o", tmp_dir] + + # fast mode + if self.use_fast_mode: + self.v("Running TotalSegmentator in fast mode ('--fast', 3mm)") + bash_command += ["--fast"] + else: + self.v("Running TotalSegmentator in default mode (1.5mm)") + + # Extract labels for left lung and right lung from total segmentator v1 output + left_lung_labels = [name for _, name in class_map["total"].items() if "left" in name and "lung" in name] + right_lung_labels = [name for _, name in class_map["total"].items() if "right" in name and "lung" in name] + + if left_lung_labels or right_lung_labels: + self.v(f"Left lung labels: {left_lung_labels}") + self.v(f"Right lung labels: {right_lung_labels}") + bash_command += ["--roi_subset"] + bash_command.extend(left_lung_labels + right_lung_labels) + + # run the model + self.v("Running",bash_command) + self.subprocess(bash_command, text=True) + + left_label_files = [os.path.join(tmp_dir, f'{i}.nii.gz') for i in left_lung_labels] + right_label_files = [os.path.join(tmp_dir, f'{i}.nii.gz') for i in right_lung_labels] + lung_left = self.n_connected(self.combine_labels(left_label_files)) + lung_right = self.n_connected(self.combine_labels(right_label_files)) + + op_data = np.zeros(lung_left.shape) + op_data[lung_left > 0] = 1 + op_data[lung_right > 0] = 2 + ref = sitk.ReadImage(in_data.abspath) + op_img = sitk.GetImageFromArray(op_data) + op_img.CopyInformation(ref) + sitk.WriteImage(op_img, out_data.abspath) \ No newline at end of file diff --git a/models/bamf_pet_ct_lung_tumor/utils/NNUnetPETCTRunner.py b/models/bamf_pet_ct_lung_tumor/utils/NNUnetPETCTRunner.py new file mode 100644 index 00000000..657f56f0 --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/utils/NNUnetPETCTRunner.py @@ -0,0 +1,117 @@ +""" +------------------------------------------------- +MHub - NNU-Net MultiModality Runner + This is a base runner for pre-trained + nnunet models +------------------------------------------------- + +------------------------------------------------- +Author: Jithendra Kumar +Email: jithendra.kumar@bamfhealth.com +------------------------------------------------- +""" +# TODO: support multi-i/o and batch processing on multiple instances + +from typing import List, Optional +import os, subprocess, shutil +import SimpleITK as sitk, numpy as np +from mhubio.core import Module, Instance, InstanceData, DataType, FileType, IO, InstanceDataCollection + +# TODO: add an optional evaluation pattern (regex) to IO.Config +nnunet_task_name_regex = r"Task[0-9]{3}_[a-zA-Z0-9_]+" + +@IO.ConfigInput('in_ct_data', 'nifti:mod=ct', the="input ct data to run nnunet on") +@IO.ConfigInput('in_pt_data', 'nifti:mod=pt', the="input pt data to run nnunet on") +@IO.Config('nnunet_task', str, None, the='nnunet task name') +@IO.Config('nnunet_model', str, None, the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)') +@IO.Config('folds', int, None, the='number of folds to run nnunet on') +@IO.Config('use_tta', bool, True, the='flag to enable test time augmentation') +@IO.Config('roi', str, None, the='roi or comma separated list of roi the nnunet segments') +class NNUnetPETCTRunner(Module): + + nnunet_task: str + nnunet_model: str + input_data_type: DataType + folds: int # TODO: support optional config attributes + use_tta: bool + roi: str + + @IO.Instance() + @IO.Input('in_ct_data', the="input ct data to run nnunet on") + @IO.Input('in_pt_data', the="input pt data to run nnunet on") + @IO.Output("out_data", 'VOLUME_001.nii.gz', 'nifti:mod=seg:model=nnunet', the="output data from nnunet") + def task(self, instance: Instance, in_ct_data: InstanceData,in_pt_data: InstanceData, out_data: InstanceData) -> None: + + # get the nnunet model to run + self.v("Running nnUNet_predict.") + self.v(f" > task: {self.nnunet_task}") + self.v(f" > model: {self.nnunet_model}") + self.v(f" > output data: {out_data.abspath}") + + # download weights if not found + # NOTE: only for testing / debugging. For productiio always provide the weights in the Docker container. + if not os.path.isdir(os.path.join(os.environ["WEIGHTS_FOLDER"], '')): + print("Downloading nnUNet model weights...") + bash_command = ["nnUNet_download_pretrained_model", self.nnunet_task] + self.subprocess(bash_command, text=True) + + # bring input data in nnunet specific format + # NOTE: only for nifti data as we hardcode the nnunet-formatted-filename (and extension) for now. + # This model expects 2 input modalities for each image + inp_dir = self.config.data.requestTempDir(label="nnunet-model-inp") + inp_file = f'VOLUME_001_0000.nii.gz' + shutil.copyfile(in_ct_data.abspath, os.path.join(inp_dir, inp_file)) + inp_file = f'VOLUME_001_0001.nii.gz' + shutil.copyfile(in_pt_data.abspath, os.path.join(inp_dir, inp_file)) + + # define output folder (temp dir) and also override environment variable for nnunet + out_dir = self.config.data.requestTempDir(label="nnunet-model-out") + os.environ['RESULTS_FOLDER'] = out_dir + + # symlink nnunet input folder to the input data with python + # create symlink in python + # NOTE: this is a workaround for the nnunet bash script that expects the input data to be in a specific folder + # structure. This is not the case for the mhub data structure. So we create a symlink to the input data + # in the nnunet input folder structure. + os.symlink(os.environ['WEIGHTS_FOLDER'], os.path.join(out_dir, 'nnUNet')) + + # NOTE: instead of running from commandline this could also be done in a pythonic way: + # `nnUNet/nnunet/inference/predict.py` - but it would require + # to set manually all the arguments that the user is not intended + # to fiddle with; so stick with the bash executable + + # construct nnunet inference command + bash_command = ["nnUNet_predict"] + bash_command += ["--input_folder", str(inp_dir)] + bash_command += ["--output_folder", str(out_dir)] + bash_command += ["--task_name", self.nnunet_task] + bash_command += ["--model", self.nnunet_model] + + # add optional arguments + if self.folds is not None: + bash_command += ["--folds", str(self.folds)] + + if not self.use_tta: + bash_command += ["--disable_tta"] + + self.v(f" > command 1: {bash_command}") + # run command + self.subprocess(bash_command, text=True) + + # output meta + meta = { + "model": "nnunet", + "nnunet_task": self.nnunet_task, + "nnunet_model": self.nnunet_model, + "roi": self.roi + } + + # get output data + out_file = f'VOLUME_001.nii.gz' + out_path = os.path.join(out_dir, out_file) + + # copy output data to instance + shutil.copyfile(out_path, out_data.abspath) + + # update meta dynamically + out_data.type.meta += meta diff --git a/models/bamf_pet_ct_lung_tumor/utils/Registration.py b/models/bamf_pet_ct_lung_tumor/utils/Registration.py new file mode 100644 index 00000000..af006863 --- /dev/null +++ b/models/bamf_pet_ct_lung_tumor/utils/Registration.py @@ -0,0 +1,55 @@ +""" +--------------------------------------------------------- +Registration Module using SimpleITK +--------------------------------------------------------- + +------------------------------------------------- +Author: Jithendra Kumar +Email: Jithendra.kumar@bamfhealth.com +------------------------------------------------- + +""" +import os +import shutil +import SimpleITK as sitk +import numpy as np +from mhubio.core import IO +from mhubio.core import Module, Instance, InstanceData + + +class Registration(Module): + + @IO.Instance() + @IO.Input('in_fixed_data', 'nifti:mod=pt', the='input fixed data') + @IO.Input('in_moving_data', 'nifti:mod=ct', the='input moving data') + @IO.Output('out_data', 'VOL000_registered.nii.gz', 'nifti:mod=ct:registered=true', the="registered ct data") + def task(self, instance: Instance, in_moving_data: InstanceData, in_fixed_data: InstanceData, out_data: InstanceData): + """ + Perform registration and resampling + """ + fixed = sitk.ReadImage(in_fixed_data.abspath, sitk.sitkFloat32) + moving = sitk.ReadImage(in_moving_data.abspath, sitk.sitkFloat32) + numberOfBins = 24 + samplingPercentage = 0.10 + R = sitk.ImageRegistrationMethod() + R.SetMetricAsMattesMutualInformation(numberOfBins) + R.SetMetricSamplingPercentage(samplingPercentage, sitk.sitkWallClock) + R.SetMetricSamplingStrategy(R.RANDOM) + R.SetOptimizerAsRegularStepGradientDescent(1.0, 0.001, 200) + R.SetInitialTransform(sitk.TranslationTransform(fixed.GetDimension())) + R.SetInterpolator(sitk.sitkLinear) + + def command_iteration(method): + print(f"{method.GetOptimizerIteration():3} = {method.GetMetricValue():10.5f}") + + R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R)) + + outTx = R.Execute(fixed, moving) + resampler = sitk.ResampleImageFilter() + resampler.SetReferenceImage(fixed) + resampler.SetInterpolator(sitk.sitkLinear) + resampler.SetDefaultPixelValue(int(np.min(sitk.GetArrayFromImage(moving)))) + resampler.SetTransform(outTx) + out = resampler.Execute(moving) + out.CopyInformation(fixed) + sitk.WriteImage(out, out_data.abspath)