nclvmass-net 1.0__py3-none-any.whl

nclvmass_net-1.0.dist-info/METADATA

@@ -0,0 +1,57 @@
+Metadata-Version: 2.4
+Name: nclvmass-net
+Version: 1.0
+Summary: An AI Model for Left Ventricular Myocardium Segmentation and Quantification from Non-Contrast Cardiac CT
+Author-email: Jia-Da Li <jiadali@ntu.edu.tw>, Tzung-Dau Wang <tdwang@ntu.edu.tw>, Wen-Jeng Lee <wenjeng@gmail.com>, Cheng-Chia Huang <r12922188@csie.ntu.edu.tw>, Chu-Song Chen <chusong@csie.ntu.edu.tw>
+License: Apache-2.0
+Project-URL: codes, https://github.com/AI-Application-and-Integration-Lab/ncLVMass-Net
+Project-URL: models, https://huggingface.co/Jia-Da/ncLVMass-Net
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Healthcare Industry
+Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: nnunetv2>=2.6.0
+Requires-Dist: huggingface-hub
+
+# ncLVMass-Net
+An AI model for left ventricular (LV) myocardium segmentation and mass quantification from non-contrast cardiac CT (NCCT) images. The model was trained with the [nnU-Net](https://github.com/mic-dkfz/nnunet) framework on 2,021 patients, all from Taiwan. It achieved an overall Dice score of 0.8827. To use it, provide an input image in NIfTI format (`.nii.gz`) and the model will output the myocardium mask, along with estimated LV mass.
+
+# Quick Install
+We recommend first installing [PyTorch](https://pytorch.org/get-started/locally/) in a Python environment (version 3.10 or later), and then installing ncLVMass-Net.
+```
+pip install nclvmass-net
+```
+
+# Usage
+```
+nclvmnet -i ct.nii.gz -o output_dir
+nclvmnet -i input_dir -o output_dir
+```
+
+# Change pretrained model
+The ncLVMass-Net pretrained model weights are automatically downloaded from [Hugging Face](https://huggingface.co/NTU-AI2Lab/ncLVMass-Net) and stored in the local Hugging Face cache. If you would like to use a different pretrained model, specify the model directory name with the `--dataset` argument and the parent directory containing the model with `--mdl-local-dir`.
+
+For example, if you have previously downloaded the [TotalSegmentator](https://github.com/wasserth/TotalSegmentator) `heartchambers_highres` model (make sure you have obtained the appropriate [license](https://backend.totalsegmentator.com/license-academic/)), you can switch to this model for comparison.
+
+```
+nclvmnet -i ct.nii.gz -o output_dir_TS --mdl-local-dir $HOME/.totalsegmentator/nnunet/results --dataset Dataset301_heart_highres_1559subj
+```
+
+For non-commercial purposes, you may download the pretrained model weights directly from Hugging Face and fine-tune them on your own data.
+```
+hf download --local-dir $HOME/path/to/ncLVMass-Net/model NTU-AI2Lab/ncLVMass-Net
+```
+
+# Further development
+The package directly invokes `nnUNetv2_predict` and `nnUNetv2_apply_postprocessing` with the specified input (`-i`), output (`-o`), and pretrained model (`--mdl-local-dir` and `--dataset`) paths, without modifying the original nnU-Net workflow. Note that the custom trainer used for other pretrained model (e.g., one with a non-default number of epochs) may not be available in the original nnU-Net codebase. For further customization or development, install the package in editable mode.
+```
+git clone https://github.com/MIC-DKFZ/nnUNet.git
+git clone https://github.com/AI-Application-and-Integration-Lab/ncLVMass-Net.git
+pip install -e ./nnUNet
+pip install -e ./ncLVMass-Net
+```

nclvmass_net-1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+nclvmnet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nclvmnet/label_mass.py,sha256=d5t44dydXIt3QibJhhjcdtwT_nufwk16ZaLFfRR80GM,4510
+nclvmnet/nclvmnet.py,sha256=xqf0bQhBAYFzmA3EkmaZjp402oQQisZOXWM82XpYAsM,12894
+nclvmass_net-1.0.dist-info/METADATA,sha256=p5XhoFQmTWtQOpQMa6fFWTcOv7cxNmyBtxVpzhFTNFQ,3590
+nclvmass_net-1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+nclvmass_net-1.0.dist-info/entry_points.txt,sha256=1zb86OfbHoJ6gTqH3HH9Kq6FLu1VacRQf-mA0pU--RE,52
+nclvmass_net-1.0.dist-info/top_level.txt,sha256=gSPHUsqsda_dGNFW8SD9yOZ6cC_qXeBTkCy7jcMHnIQ,9
+nclvmass_net-1.0.dist-info/RECORD,,

nclvmass_net-1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nclvmass_net-1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+nclvmnet = nclvmnet.nclvmnet:main

nclvmass_net-1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+nclvmnet

nclvmnet/label_mass.py

@@ -0,0 +1,112 @@
+import argparse
+import os
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+NIFTI_EXTS = (".nii.gz", ".nii")
+SITK_EXTS = (".nrrd", ".nhdr", ".mha", ".mhd")
+ALL_EXTS = NIFTI_EXTS + SITK_EXTS
+
+NIFTI_UNIT_TO_MM = {"unknown": 1.0, "meter": 1000.0, "mm": 1.0, "micron": 0.001}
+
+# --------------------------------------------------------------------------------
+# Format-specific readers: each returns (3D numpy array, spacing in mm as a tuple)
+# --------------------------------------------------------------------------------
+def read_nifti(entry: os.DirEntry):
+    import nibabel as nib
+
+    img = nib.load(entry.path)
+    data = np.asanyarray(img.dataobj)
+    zooms = img.header.get_zooms()[: data.ndim]
+    unit = img.header.get_xyzt_units()[0]
+    factor = NIFTI_UNIT_TO_MM.get(unit, 1.0)
+    if unit not in NIFTI_UNIT_TO_MM:
+        print(f"[Warning] {entry.name}: unknow unit '{unit}', assuming mm", file=sys.stderr)
+    spacing = tuple(float(z) * factor for z in zooms)
+    return data, spacing
+
+def read_sitk(entry: os.DirEntry):
+    import SimpleITK as sitk
+
+    img = sitk.ReadImage(entry.path)
+    data = sitk.GetArrayFromImage(img)  # array order is (z, y, x)
+    spacing = tuple(float(s) for s in img.GetSpacing())  # spacing order is (x, y, z), in mm
+    return data, spacing
+
+def load_segmentation(entry: os.DirEntry):
+    name = entry.name.lower()
+    if name.endswith(NIFTI_EXTS):
+        data, spacing = read_nifti(entry)
+    elif name.endswith(SITK_EXTS):
+        data, spacing = read_sitk(entry)
+    else:
+        raise ValueError(f"Unsupported file format: {entry.name}")
+
+    # Handle non-single-volume data (e.g. 4D, or an extra channel dimension)
+    data = np.squeeze(data)
+    if data.ndim != 3:
+        raise ValueError(f"{entry.name}: data is {data.ndim}D, only 3D is supported (shape={data.shape})")
+    if len(spacing) != 3:
+        raise ValueError(f"{entry.name}: got {len(spacing)} spacing values, expected 3 ({spacing})")
+    return data, spacing
+
+# ------------------
+# Volume computation
+# ------------------
+def compute_volumes(data: np.ndarray, spacing_mm, target_label:int):
+    """Return {label: (voxel_count, volume_cm3)}"""
+    if not np.issubdtype(data.dtype, np.integer):
+        rounded = np.rint(data)
+        if not np.allclose(data, rounded):
+            raise ValueError(
+                "Data contains non-integer values; this does not look like a label map "
+                "(is it a softmax / probability output?)"
+            )
+        data = rounded.astype(np.int64)
+
+    voxel_mm3 = float(np.prod(spacing_mm))
+    result = np.sum(data==target_label) * voxel_mm3 / 1000.0  # mm^3 -> cm^3
+    return result
+
+def get_label(dataset_json, label_name):
+    """Read {int: name} from an nnU-Net dataset.json."""
+    with open(dataset_json, 'rt') as f:
+        dataset = json.load(f)
+    labels = dataset.get('labels', {})
+    label = None
+    for k, v in labels.items():
+        if isinstance(k, str) and label_name in k.lower():
+            label = int(v)
+        elif isinstance(v, str) and label_name in v.lower():
+            label = int(k)
+    return label
+
+def label_mass(label_name:str):
+    parser = argparse.ArgumentParser(description="Compute specific label volume (cm^3) and mass (g) from nnU-Net segmentations")
+    parser.add_argument("--input", type=str, required=True, help="The directory path which conaining mask files output from nnU-Net framework")
+    parser.add_argument("--dataset-json", type=str, required=True, help="nnU-Net dataset.json, used to display label names")
+    parser.add_argument("--rho", type=float, default=1.055, help="Myocardium density in g/cm^3. Default is 1.055 g/cm^3.")
+    args = parser.parse_args()
+
+    out = 'density (g/cm^3): ' + str(args.rho) + '\n'
+    for entry in os.scandir(args.input):
+        if entry.is_file() and entry.name.lower().endswith(ALL_EXTS):
+            data, spacing = load_segmentation(entry)
+            label = get_label(args.dataset_json, label_name)
+            if label is None:
+                print(f"[Warning] Label '{label_name}' not found in {args.dataset_json}")
+                continue
+            volume_cm3 = compute_volumes(data, spacing, label)
+            mass_g = volume_cm3 * args.rho
+            out += f"{entry.name}: volume={volume_cm3:.3f} cm^3, mass={mass_g:.3f} g\n"
+
+    txt_path = os.path.join(args.input, f'{label_name}_mass.txt')
+    with open(txt_path, 'wt') as f:
+        f.write(out)
+
+
+if __name__ == "__main__":
+    label_mass("myocardium")

nclvmnet/nclvmnet.py

@@ -0,0 +1,219 @@
+import os
+import argparse
+import tempfile
+import json
+import sys
+import shutil
+from huggingface_hub import snapshot_download, try_to_load_from_cache
+from nnunetv2.inference.predict_from_raw_data import predict_entry_point
+from nnunetv2.postprocessing.remove_connected_components import entry_point_apply_postprocessing
+from nclvmnet.label_mass import label_mass
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ncLVMass-Net: A tool for estimating left ventricular mass from non-contrast CT images using nnU-Net.")
+    parser.add_argument('-i', '--input', type=str, required=True, help="A CT nifti image path or a folder containing CT nifti images.")
+    parser.add_argument('-o', '--output', type=str, required=True, help='Output folder path.')
+    parser.add_argument('--mdl-local-dir', type=str, default=None, help='Local directory of the model folder.')
+    parser.add_argument('--hf-token', type=str, default=None, help='Hugging Face token for accelerate download process.')
+    parser.add_argument('-d', '--dataset', type=str, default=None, help='Select a dataset in the model folder to predict. If not specified, the first finding dataset will be selected.')
+    parser.add_argument('-tr', '--trainer', type=str, default=None, help='What nnU-Net trainer class was used for training? If not specified, the first finding folder will be selected.')
+    parser.add_argument('-p', '--plans', type=str, default=None, help='Plans identifier. If not specified, the first finding folder will be selected.')
+    parser.add_argument('-c', '--config', type=str, default=None, help='nnU-Net configuration (2d, 3d_fullres, 3d_lowres, 3d_cascade_fullres). If not specified, the first finding folder will be selected.')
+    parser.add_argument('-f', '--fold', nargs='+', type=str, default=['0','1','2','3','4'], help='Which fold(s) to use for prediction. Default is depends on first finding "crossval_results_folds_x_x_x..." folder. If not exist, select "all" if exist, otherwise select the first fold.')
+    parser.add_argument('-chk', '--checkpoint', type=str, default='checkpoint_best.pth', help='Checkpoint file name. Should be checkpoint_best.pth (default) or checkpoint_final.pth.')
+    parser.add_argument('--myo-density', type=float, default=1.055, help='Myocardium density in g/cm^3. Default is 1.055 g/cm^3.')
+    args = parser.parse_args()
+
+    # ---------------------------------
+    # verify and update input arguments
+    # ---------------------------------
+    if args.mdl_local_dir is None:
+        # Download repository from huggingface and get the local path of the model folder
+        repo_id = "Jia-Da/ncLVMass-Net"
+        # repo_id = "NTU-AI2Lab/ncLVMass-Net"
+        readme_path = try_to_load_from_cache(repo_id=repo_id, filename="README.md")
+        if readme_path is not None:
+            args.mdl_local_dir = os.path.dirname(readme_path)
+            print('ncLVMass-Net model weights already downloaded.')
+        else:
+            if args.hf_token is None:
+                args.mdl_local_dir = snapshot_download(repo_id=repo_id)
+            else:
+                args.mdl_local_dir = snapshot_download(repo_id=repo_id, token=args.hf_token)
+    else:
+        if not os.path.isdir(args.mdl_local_dir):
+            raise ValueError(f"The specified model local directory does not exist: {args.mdl_local_dir}")
+
+    # -------------------
+    # find Dataset folder
+    # -------------------
+    if args.dataset is None:
+        args.dataset = next(f.name for f in os.scandir(args.mdl_local_dir) if f.is_dir() and f.name.startswith("Dataset"))
+    else:
+        if not os.path.isdir(os.path.join(args.mdl_local_dir, args.dataset)):
+            if args.dataset.isdecimal():
+                args.dataset = next(f.name for f in os.scandir(args.mdl_local_dir) if f.is_dir() and f.name.startswith(f"Dataset{args.dataset}"))
+            else:
+                args.dataset = next(f.name for f in os.scandir(args.mdl_local_dir) if f.is_dir() and f.name.startswith("Dataset") and args.dataset in f.name)
+
+    if args.trainer is None or \
+       args.plans is None or \
+       args.config is None or \
+       not os.path.isdir(os.path.join(args.mdl_local_dir, args.dataset, f'{args.trainer}__{args.plans}__{args.config}')):
+        # find plans, trainer, config
+        for f in os.scandir(os.path.join(args.mdl_local_dir, args.dataset)):
+            if f.is_dir() and '__' in f.name:
+                parts = f.name.split('__')
+                if len(parts) == 3:
+                    args.trainer, args.plans, args.config = parts
+                    break
+
+    subfolder = os.path.join(args.mdl_local_dir, args.dataset, f'{args.trainer}__{args.plans}__{args.config}')
+    if os.path.isdir(os.path.join(subfolder, 'crossval_results_folds_' + '_'.join(args.fold))):
+        postprocess = True
+        for f in args.fold:
+            assert os.path.isdir(os.path.join(subfolder, f'fold_{f}')), 'Fold folder does not exist: ' + os.path.join(subfolder, f'fold_{f}')
+            if not os.path.isfile(os.path.join(subfolder, f'fold_{f}', args.checkpoint)):
+                if f in {'0', 'all'}:
+                    for chk in ('checkpoint_best.pth', 'checkpoint_final.pth'):
+                        if os.path.isfile(os.path.join(subfolder, f'fold_{f}', chk)):
+                            args.checkpoint = chk
+                            break
+                    assert os.path.isfile(os.path.join(subfolder, f'fold_{f}', args.checkpoint)), 'Checkpoint file does not exist: ' + os.path.join(subfolder, f'fold_{f}', args.checkpoint)
+                else:
+                    raise FileNotFoundError('Checkpoint file does not exist: ' + os.path.join(subfolder, f'fold_{f}', args.checkpoint))
+    else:
+        crf = next((f.name for f in os.scandir(subfolder) if f.is_dir() and f.name.startswith('crossval_results_folds_')), None)
+        if crf is not None:
+            if crf.split('_')[3:] != args.fold:
+                print(f'Warning: The specified folds {args.fold} do not match the folds in the cross-validation results folder "{crf}". Using folds from the folder instead.')
+                args.fold = crf.split('_')[3:]
+            postprocess = True
+            for f in args.fold:
+                assert os.path.isdir(os.path.join(subfolder, f'fold_{f}')), 'Fold folder does not exist: ' + os.path.join(subfolder, f'fold_{f}')
+                if not os.path.isfile(os.path.join(subfolder, f'fold_{f}', args.checkpoint)):
+                    if f in {'0', 'all'}:
+                        for chk in ('checkpoint_best.pth', 'checkpoint_final.pth'):
+                            if os.path.isfile(os.path.join(subfolder, f'fold_{f}', chk)):
+                                args.checkpoint = chk
+                                break
+                        assert os.path.isfile(os.path.join(subfolder, f'fold_{f}', args.checkpoint)), 'Checkpoint file does not exist: ' + os.path.join(subfolder, f'fold_{f}', args.checkpoint)
+                    else:
+                        raise FileNotFoundError('Checkpoint file does not exist: ' + os.path.join(subfolder, f'fold_{f}', args.checkpoint))
+        else:
+            postprocess = False
+            fold_folder = next((f.name for f in os.scandir(subfolder) if f.is_dir() and f.name.startswith('fold_')), None)
+            if fold_folder is not None:
+                if args.fold != [fold_folder.split('_')[1]]:
+                    print(f'Warning: The specified folds {args.fold} do not match the fold in the folder "{fold_folder}". Using fold from the folder instead.')
+                    args.fold = [fold_folder.split('_')[1]]
+                if not os.path.isfile(os.path.join(subfolder, fold_folder, args.checkpoint)):
+                    for chk in ('checkpoint_best.pth', 'checkpoint_final.pth'):
+                        if os.path.isfile(os.path.join(subfolder, fold_folder, chk)):
+                            args.checkpoint = chk
+                            break
+                    assert os.path.isfile(os.path.join(subfolder, fold_folder, args.checkpoint)), 'Checkpoint file does not exist: ' + os.path.join(subfolder, fold_folder, args.checkpoint)
+            else:
+                raise FileNotFoundError('No fold folder found in: ' + subfolder)
+   
+    # ---------------------------------------
+    # check input image file name and output folder
+    # ---------------------------------------
+    if os.path.exists(args.input):
+        args.input = os.path.abspath(args.input)
+    else:
+        raise FileNotFoundError(f"The specified input path does not exist: {args.input}")
+
+    dataset_json_path = os.path.join(subfolder, 'dataset.json')
+    if not os.path.isfile(dataset_json_path):
+        raise FileNotFoundError(f"dataset.json file does not exist in the model folder: {dataset_json_path}")
+    with open(dataset_json_path, 'rt') as f:
+        dataset_json = json.load(f)
+
+    # find the CT channel index
+    ct_channel = None
+    for k, v in dataset_json['channel_names'].items():
+        if v.upper()=='CT':
+            ct_channel = k.zfill(4)
+            break
+    assert ct_channel is not None, 'CT channel not found in dataset.json'
+
+    # build a temporary folder and copy input images in it, filename adding the CT channel suffix if necessary
+    if os.path.isfile(args.input):
+        src_img_dir = os.path.dirname(args.input)
+        src_img_name = [os.path.basename(args.input)]
+    else:
+        src_img_dir = args.input
+        src_img_name = [f.name for f in os.scandir(src_img_dir) if f.is_file() and f.name.endswith(dataset_json['file_ending'])]
+    if len(src_img_name) == 0:
+        raise FileNotFoundError(f"No input image files found in the specified path: {args.input}")
+    extL = len(dataset_json['file_ending'])
+    tmp_img_dir = tempfile.mkdtemp(dir=src_img_dir)
+    tmp_img_name = []
+    for src_name in src_img_name:
+        src = os.path.join(src_img_dir, src_name)
+        if src_name.endswith(f'_{ct_channel}{dataset_json["file_ending"]}'):
+            tmp_img_name.append(src_name)
+        else:
+            tmp_img_name.append(f'{src_name[:-extL]}_{ct_channel}{dataset_json["file_ending"]}') # add channel
+        dst = os.path.join(tmp_img_dir, tmp_img_name[-1])
+        shutil.copyfile(src, dst)
+
+    # prepare output folder
+    if os.path.exists(args.output):
+        if not os.path.isdir(args.output):
+            raise NotADirectoryError(f"The specified output path is not a directory: {args.output}")
+    else:
+        os.makedirs(args.output)
+     
+    # ----------------------------------------
+    # create environment variables for nnU-Net
+    # ----------------------------------------
+    if os.environ.get('nnUNet_raw') is None:
+        os.environ['nnUNet_raw'] = '<path_to_nnUNet_raw>'
+    if os.environ.get('nnUNet_preprocessed') is None:
+        os.environ['nnUNet_preprocessed'] = '<path_to_nnUNet_preprocessed>'
+    # overwrite nnUNet_results to the model folder
+    os.environ['nnUNet_results'] = args.mdl_local_dir
+
+    # ======================
+    # run nnU-Net prediction
+    # ======================
+    if postprocess:
+        print('Starting prediction...')
+        output_predict = os.path.join(tmp_img_dir, 'predicted')
+        sys.argv = ["predict_entry_point", "-i", tmp_img_dir, "-o", output_predict, "-d", args.dataset, "-tr", args.trainer, "-p", args.plans, "-c", args.config, "-f"] + args.fold + ["-chk", args.checkpoint]
+        predict_entry_point()
+        
+        # post-processing
+        print('Applying post-processing...')
+        pkl_file = os.path.join(subfolder, 'crossval_results_folds_' + '_'.join(args.fold), 'postprocessing.pkl')
+        plans_json_file = os.path.join(subfolder, 'plans.json')
+        sys.argv = ["entry_point_apply_postprocessing", "-i", output_predict, "-o", args.output, "-pp_pkl_file", pkl_file, "-plans_json", plans_json_file, "-dataset_json", dataset_json_path]
+        entry_point_apply_postprocessing()
+
+        # Move other files (e.g., *.json) in the output_predict folder to the final output folder
+        for f in os.scandir(output_predict):
+            if f.is_file() and not f.name.endswith(dataset_json['file_ending']):
+                os.rename(f.path, os.path.join(args.output, f.name))
+    else:
+        print('Starting prediction...')
+        sys.argv = ["predict_entry_point", "-i", tmp_img_dir, "-o", args.output, "-d", args.dataset, "-tr", args.trainer, "-p", args.plans, "-c", args.config, "-f"] + args.fold + ["-chk", args.checkpoint]
+        predict_entry_point()
+
+    # remove image files and temp directory after prediction
+    shutil.rmtree(tmp_img_dir)
+
+    # ----------------
+    # estimate LV mass
+    # ----------------
+    print('Estimating LV mass...')
+    sys.argv = ['label_mass', '--input', args.output, '--dataset-json', dataset_json_path, '--rho', str(args.myo_density)]
+    label_mass('myocardium')
+    print('All done!')
+
+
+
+if __name__ == "__main__":
+    main()