hyperresashs 1.0.0__py3-none-any.whl

hyperresashs/__init__.py

@@ -0,0 +1,25 @@
+import sys
+import os
+from importlib.metadata import version, PackageNotFoundError
+
+# Check for any already-cached nnunetv2 modules
+to_remove = [key for key in sys.modules if key == "nnunetv2" or key.startswith("nnunetv2.")]
+for key in to_remove:
+    print(f'hyperresashs found existing nnunetv2 module {key} in sys.modules. Removing it to avoid conflicts with hyperresashs submodule nnunetv2.')
+    del sys.modules[key]
+
+# Insert the submodule path at the front so it takes priority
+_submodule_path = os.path.join(os.path.dirname(__file__), "submodules", "nnUNet")
+sys.path.insert(0, _submodule_path)
+
+# Insert the submodule path at the front so it takes priority
+_submodule_path_inr = os.path.join(os.path.dirname(__file__), "submodules", "multi_contrast_inr")
+sys.path.insert(0, _submodule_path_inr)
+
+# Get the version from python importlib
+try:
+    # Replace 'your_package_name' with the actual name of your package
+    __version__ = version("hyperresashs") 
+except PackageNotFoundError:
+    # Handle cases where the package isn't installed in the environment (e.g., running from source without proper installation)
+    __version__ = "unknown"

hyperresashs/__main__.py

@@ -0,0 +1,2 @@
+from .ashs_cli import main
+main()

hyperresashs/ashs_cli.py

@@ -0,0 +1,511 @@
+from .ashs_inference import HyperASHSInference
+from .utils.huggingface import hf_disable_ssl_verification, hf_read_yaml, torch_hub_disable_ssl_verification
+from .ashs_training import HyperASHSTraining
+from .utils.tool import copy_or_link_file
+from . import __version__
+import argparse
+import huggingface_hub as hf
+import pandas as pd
+import re
+import pprint
+import json
+import textwrap
+import os
+import shutil
+import yaml
+import sys
+from typing import Dict, Any
+from importlib.resources import files
+
+"""
+Original ASHS command line help message for reference:
+		required options:
+		  -a dir            Location of the atlas directory. Can be a full pathname or a
+		                    relative directory name under ASHS_ROOT/data directory. 
+		  -g image          Filename of 3D (g)radient echo MRI (ASHS_MPRAGE, T1w)
+		  -f image          Filename of 2D focal (f)ast spin echo MRI (ASHS_TSE, T2w)
+		  -w path           Working/output directory
+		  
+		optional:
+		  -d                Enable debugging
+		  -h                Print help
+		  -s integer        Run only one stage (see below); also accepts range (e.g. -s 1-3)
+		  -N                No overriding of registration results. If a result from an earlier run
+		                    exists, don't run greedy again.
+		  -G                Use template brain mask in T1 template rigid registratian
+		  -T                Tidy mode. Cleans up files once they are unneeded. The -N option will
+		                    have no effect in tidy mode, because intermediate results will be erased.
+		  -I string         Subject ID (for stats output). Defaults to last word of working dir.
+		  -V                Display version information and exit
+		  -C file           Configuration file. If not passed, uses $ASHS_ROOT/bin/ashs_config.sh
+		  -Q                Use Sun Grid Engine (SGE) to schedule sub-tasks in each stage. 
+		                    By default, the whole ashs_main job runs in a single process. 
+		                    If you are doing a lot of segmentations and have SGE, it is better to 
+		                    run each segmentation (ashs_main) in a separate SGE job, rather than use the -q flag. 
+		                    The -q flag is best for when you have only a few segmentations and want them to run fast.
+		  -P                Use GNU parallel to run on multiple cores on the local machine. You need to
+		                    have GNU parallel installed.
+		  -S                Use SLURM instead of SGE, LSF or GNU parallel
+		  -l                Use LSF instead of SGE, SLURM or GNU parallel
+		  -q OPTS           Pass in additional options to SGE/SLURM/LSF/GNU Parallel. If -S -B or -P not specified
+		                    turns on SGE. 
+		  -z script         Provide a path to an executable script that will be used to retrieve SGE, LSF, SLURM or
+		                    GNU parallel options for different stages of ASHS. Takes precendence over -q
+		  -r files          Compare segmentation results with a reference segmentation. The parameter
+		                    files should consist of two nifti files in quotation marks:
+		                    
+		                      -r "ref_seg_left.nii.gz ref_seg_right.nii.gz"
+		                    
+		                    The results will include overlap calculations between different
+		                    stages of the segmentation and the reference segmentation. Note that the
+		                    comparison takes into account the heuristic rules specified in the altas, so
+		                    it is not as simple as computing dice overlaps between the reference seg
+		                    and the ASHS segs.
+		  -m file           Provide the .mat file for the transform between the T1w and T2w image. The file
+		                    is in the format used by ITK-SNAP and C3D and should be generated by performing
+		                    registration with the T2w MRI as fixed image and the T1w MRI as moving. 
+		                    By default, the mat file is used as a hint to initialize rigid T2/T1
+		                    but this can be modified with the -M flag.
+		  -M                The mat file provided with -m is used as the final T2/T1 registration.
+		                    ASHS will not attempt to run registration between T2 and T2.
+		  -t threads        Specify number of parallel threads the greedy runs
+		  -H                Tell ASHS to use external hooks for reporting progress, errors, and warnings.
+		                    The environment variables ASHS_HOOK_SCRIPT must be set to point to the appropriate
+		                    script. For an example script with comments, see ashs_default_hook.sh
+		                    The purpose of the hook is to allow intermediary systems (e.g. XNAT) 
+		                    to monitor ASHS performance. An optional ASHS_HOOK_DATA variable can be set
+		  -B                Do not perform the bootstrapping step, and use the output of the initial joint
+		                    label fusion (in multiatlas directory) as the final output.
+		                    and will be forwarded to the script
+"""
+
+_ashs_naming_scheme = """
+# ------- input -------
+# input image names
+t1_native_img: "mprage.nii.gz"
+t2_native_img: "tse.nii.gz"
+
+# input template names
+template: "template.nii.gz"
+left_roi_file: "left_round_in_global_space_larger.nii.gz"
+right_roi_file: "right_round_in_global_space_larger.nii.gz"
+
+# ------- output -------
+# registration intermediate
+t1_name_after_triming_neck: "mprage_necktrim.nii.gz"
+affine_matrix: "t1_to_template_affine_inv.mat"
+t1_whole_img: "mprage_to_tse_warped.nii.gz"
+t2_padded_img: "tse_padded.nii.gz"
+
+# registration output
+template_to_3tt1: "template_to_mprage_warped.nii.gz"
+global_roi_in_3tt1_XYZ: "template_roi_XYZ_to_mprage_warped.nii.gz"
+
+# preprocessing
+inr_primary: "primary.nii.gz"
+inr_secondary: "secondary.nii.gz"
+inr_seg: "primary_seg.nii.gz"
+seg: "input_primary_seg.nii.gz"
+hyper_primary: "hyper_primary.nii.gz"
+hyper_secondary: "hyper_secondary.nii.gz"
+reg_mat: "auxiluary_to_primary.mat"
+hyper_secondary_after_registertion: "auxiluary_to_primary_registered.nii.gz"
+hyper_primary_seg: "inr_hyper_primary_seg.nii.gz"
+"""
+
+
+def main():
+    
+    parser = argparse.ArgumentParser(
+        prog='hrashs',
+        description='HyperResASHS: High-Resolution Automatic Segmentation of Hippocampal Subfields',
+        epilog="""HyperResASHS (C) 2026 by Yue Li, Paul Yushkevich, and UPenn Patch Lab.
+                    Citation: https://doi.org/10.48550/arXiv.2508.17171
+                    GitHub:   https://github.com/liyue3780/HyperResASHS.git""")
+    
+    # Create subparsers for different commands
+    subparsers = parser.add_subparsers(dest='command', required=True, help='Available commands')
+    
+    # List atlases subcommand
+    list_parser = subparsers.add_parser('list', help='List available HyperResASHS atlases')
+    list_parser.add_argument('-l', '--long', action='store_true', help='Show detailed information about each atlas')
+    
+    # Describe atlas subcommand
+    describe_parser = subparsers.add_parser('desc', help='Describe a specific HyperResASHS atlas')
+    describe_parser.add_argument('atlas', type=str, help='Name of the atlas to describe or path to atlas config file')
+    
+    # Run segmentation subcommand
+    run_parser = subparsers.add_parser('run', help='Run HyperResASHS segmentation pipeline')
+    run_parser.add_argument('-a', '--atlas', type=str, required=True, 
+                           help='Name of the atlas to use or path to atlas config file')
+    run_parser.add_argument('-g', '--t1', type=str, required=True, 
+                           help='Path to T1-weighted image')
+    run_parser.add_argument('-f', '--t2', type=str, required=True, 
+                           help='Path to T2-weighted image')
+    
+    # Training command
+    train_parser = subparsers.add_parser('train', help='Train a HyperResASHS model on a dataset')
+    train_parser.add_argument('-c', '--config', type=str, required=True, 
+                              help='Path to training configuration YAML file. See documentation for expected format.')
+    train_parser.add_argument('-m', '--manifest', type=str, required=True, 
+                              help='Path to manifest CSV file describing the training dataset. See documentation for expected format.')
+    train_parser.add_argument('-l', '--labels', type=str, required=True, 
+                              help='Path to ITK-SNAP label description file. See documentation for expected format.')
+    train_parser.add_argument('-x', '--xval', type=str, default=None,
+                              help='Cross-validation fold specification. See documentation for expected format.')
+    train_parser.add_argument('-R', '--inr-random-seed', type=int, default=None, 
+                              help='Specify random seed for the INR optimization; use to rerun failed INR experiments.')
+    train_parser.add_argument('--inr-batch-size', type=int, default=None, 
+                              help='Specify batch size for the INR optimization. Default: 10000')
+    
+    train_parser.add_argument('-s', '--stage', type=str, default='',
+                              help='''
+                              Run one or more selected stages of the training pipeline.
+                              You can specify single stage as -s 1, or a range of stages as -s 1-3.
+                              The stages are as follows:
+                                1: Preprocessing (neck trim, registration, patch extraction)
+                                2: INR training
+                                3: nnU-Net preparation (resampling, cropping, and formatting for nnU-Net)
+                                4: nnU-Net training                                
+                              ''')
+    train_parser.add_argument('-F', '--filter', type=str, metavar='REGEX' ,default=None,
+                              help='''Restrict execution to image(s) that match specified regular expression. 
+                              For stage 1 (preprocessing), REGEX will be matched to subject id and date.
+                              For stage 2 (INR training), REGEX will be matched to subject id, date, and side.
+                              For stage 3 (nnU-Net preparation), this is ignored.
+                              For stage 4 (nnU-Net training), REGEX is matched to fold number (0-4)
+                              ''')
+
+
+    # Add common arguments for run and train subcommands
+    for p in [run_parser, train_parser]:
+        p.add_argument('-w', '--workdir', type=str, required=True, 
+                       help='Path to working directory')
+        p.add_argument('-N', '--no-overwrite', action='store_true', 
+                       help='Do not overwrite existing results')
+        p.add_argument('-t', '--threads', type=int, default=1, 
+                       help='Number of parallel threads to use for segmentation [default: 1]')
+        p.add_argument('--device', type=str, default='auto', 
+                       help='''Device to use for segmentation (e.g. "cuda" or "cpu"). Default: auto-detect
+                       To select specific GPU(s), use the CUDA_VISIBLE_DEVICES environment variable, e.g.:
+                         CUDA_VISIBLE_DEVICES=0 hrashs run ...''')
+        p.add_argument('-L', '--no-links', action='store_true', 
+                       help='Do not create symlinks in the working directory; copy files instead. By default, symlinks are created to the input T1 and T2 images in the working directory')
+        p.add_argument('-T', '--tidy', action='store_true',
+                       help='Tidy mode. Reduce the number of intermediate files generated.')
+
+    # Add -k flag for all relevant sub parsers
+    for p in [list_parser, run_parser, describe_parser, train_parser]:
+        p.add_argument('-k', '--disable-ssl-verification', action='store_true', 
+                       help='Disable SSL verification for Hugging Face Hub access')
+    
+    args = parser.parse_args()
+    
+    # Disable SSL verification if requested
+    if args.disable_ssl_verification:
+        hf_disable_ssl_verification()
+        torch_hub_disable_ssl_verification()
+    
+    # Handle commands
+    if args.command == 'list':
+        list_atlases(args)
+    elif args.command == 'desc':
+        describe_atlas(args)
+    elif args.command == 'run':
+        run_segmentation(args)
+    elif args.command == 'train':
+        return run_training(args)
+    else:
+        raise ValueError(f"Unknown command: {args.command}")
+        
+    return 0
+
+    
+def get_atlas_listing(match=None):
+    # Fetch upennpatchlab/hyperresashs_atlas_directory file active_atlases.yaml from HF:    
+    # Read the list of available atlases from the Hugging Face Hub
+    repo_list = hf_read_yaml(
+        repo_id="upennpatchlab/hyperresashs_atlas_directory", 
+        filename="active_atlases.yaml")        
+    d = {}
+    for repo in repo_list:
+        repo_config = hf_read_yaml(repo_id=repo, filename=f"atlas.yaml")
+        if match and not re.match(match, repo_config['metadata'].get('id',''), re.IGNORECASE):
+            continue
+        repo_config['metadata']['json'] = json.dumps(repo_config)
+        repo_config['metadata']['repo'] = repo
+        for k, v in repo_config['metadata'].items():
+            if k not in d:
+                d[k] = []
+            d[k].append(v)
+            
+    df = pd.DataFrame(d)
+    return df  
+    
+    
+def print_atlas_listing(long=False, match=None):
+    # Get the atlas listing as a DataFrame
+    df = get_atlas_listing(match=match)
+    df = df[[c for c in df.columns if c != 'json']]  # Exclude the json column from display
+    df.columns = [c.capitalize().replace('_',' ') for c in df.columns]  
+    
+    # Print the atlas directory in a nice format
+    if long:
+        for _, row in df.iterrows():
+            pprint.pp(row.to_dict())
+    else:
+        print(df[['Id','Name']].to_string(index=False))            
+            
+
+def list_atlases(args):
+    """List available atlases."""
+    try: 
+        print_atlas_listing(long=args.long)
+    except Exception as e:
+        print(f"Error fetching atlas directory from Hugging Face Hub: {e}")
+        return
+    
+
+def describe_atlas(args):
+    """Describe a specific atlas."""
+    # Read the list of available atlases from the Hugging Face Hub
+    try: 
+        print_atlas_listing(long=True, match=args.atlas)
+    except Exception as e:
+        print(f"Error fetching atlas directory from Hugging Face Hub: {e}")
+        return
+    
+    
+def print_header(metadata):
+
+    main_citation = """\
+        Li, Y., Khandelwal, P., Jena, R., Xie, L., Duong, M., Denning, A.E., Brown, C.A.,
+        Wisse, L.E., Das, S.R., Wolk, D.A. and Yushkevich, P.A., 2025.
+        Achieving detailed medial temporal lobe segmentation with upsampled
+        isotropic training from implicit neural representation.
+        arXiv preprint arXiv:2508.17171."""
+    nnunet_citation = """\
+        Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). 
+        nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. 
+        Nature methods, 18(2), 203-211."""
+    citations = [main_citation, nnunet_citation] + metadata.get('citations', [])
+
+    print('' + '='*80)
+    print(f"HyperResASHS {__version__} using atlas {metadata['id']} {metadata['version']}")
+    print(f"     - {metadata['name']}")
+    print()
+    print(f"Please cite the following papers if you use this atlas in your research:")
+    
+    for i,ref in enumerate(citations):
+        print(f"\n[{i+1}]  {textwrap.fill(textwrap.dedent(ref), width=74, subsequent_indent=' ' * 5)}")
+    print('' + '='*80)
+    
+    
+class Logger(object):
+    
+    class Stream(object):
+        def __init__(self, original_stream, log_file):
+            self.original_stream = original_stream
+            self.log_file = log_file
+
+        def write(self, message):
+            self.original_stream.write(message) # Write to console
+            self.log_file.write(message)      # Write to file
+
+        def flush(self):
+            # This flush method is important for proper buffering behavior
+            self.original_stream.flush()
+            self.log_file.flush()
+    
+    def __init__(self, filename):
+        self.log_file = open(filename, "w")
+        self.out = self.Stream(sys.stdout, self.log_file)
+        self.err = self.Stream(sys.stderr, self.log_file)
+        sys.stdout = self.out # Redirect all print statements to this object
+        sys.stderr = self.err # Redirect all error messages to this object
+        
+    def close(self):
+        self.log_file.close()
+        sys.stdout = self.out.original_stream # Restore original stdout
+        sys.stderr = self.err.original_stream # Restore original stderr
+        
+    def __del__(self):
+        self.close()
+        
+        
+def _fetch_template(atlas_config : Dict[str,Any]) -> str:
+    # Fetch the template from Hugging Face Hub if specified in the atlas config
+    template_hf = atlas_config.get('template', {}).get('hf', None)
+    if template_hf:
+        template_path = hf.snapshot_download(template_hf)
+    else:
+        template_path = atlas_config.get('template', {}).get('local', None)
+        if template_path is None:
+            raise ValueError("No template specified in atlas configuration. Please check the atlas configuration and try again.")
+        
+    return template_path
+
+
+def _setup_config(atlas_config : Dict[str,Any], args: argparse.Namespace, atlas_local_path:str, training:bool=False):
+    
+    # Set up the atlas configuration the way Yue's code expects it
+    config_src = atlas_config['config']
+    config_src['TEST_PATH'] = args.workdir
+    config_src['TEMPLATE_PATH'] = os.path.join(_fetch_template(atlas_config), 'hyperashs-template')
+    config_src['ATLAS_PATH'] = atlas_local_path
+    config_src['ITKSNAP_LABEL_FILE'] = os.path.join(atlas_local_path, 'itksnap_labels.txt')
+    config_src['GREEDY_NUM_THREADS'] = args.threads
+    config_src['NNUNET_NUM_THREADS'] = args.threads
+    config_src['FILE_NAME_CONFIG'] = os.path.join(args.workdir, 'config', 'ashs_filename_scheme.yaml')
+    
+    # Write the config to the working directory
+    os.makedirs(os.path.join(args.workdir, 'config'), exist_ok=True)
+    fn_config_yaml = os.path.join(args.workdir, 'config', f'configtest_{config_src["EXP_NUM"]}_{config_src["MODEL_NAME"]}.yaml')
+    with open(fn_config_yaml, 'wt') as f:
+        yaml.dump(config_src, f)
+        
+    # Write the preferred naming scheme to the working directory
+    fn_scheme = files('hyperresashs').joinpath('config_templates/ashs_filename_scheme.yaml')
+    with fn_scheme.open('r') as f:
+        _ashs_naming_scheme = f.read()
+    with open(config_src['FILE_NAME_CONFIG'], 'wt') as f:
+        f.write(_ashs_naming_scheme)
+    
+    # Create the config in the way that Yue's code expects it
+    config = yaml.safe_load(open(fn_config_yaml, 'r'))
+    
+    return config
+
+        
+
+def run_segmentation(args):
+    """Run the segmentation pipeline."""
+    
+    # Create a logger
+    os.makedirs(os.path.join(args.workdir, 'logs'), exist_ok=True)
+    logger = Logger(os.path.join(os.path.join(args.workdir, 'logs'), 'hyperashs_log.txt'))
+    
+    # First, the atlas may be either a local path to the atlas or a huggingface hub id.
+    atlas_local_path = None
+    if os.path.isdir(args.atlas) and os.path.exists(args.atlas):
+        # If it's a local path, we assume it's a config file and we load the config directly from it
+        with open(os.path.join(args.atlas, 'atlas.yaml'), 'r') as f:
+            atlas_config = yaml.safe_load(f)
+            atlas_local_path = args.atlas
+
+        # Print the header with atlas information and citations
+        print_header(atlas_config['metadata'])
+    else:
+        # Fetch the atlas the user wants
+        df_meta = get_atlas_listing(match=args.atlas)
+        if len(df_meta) == 0:
+            raise ValueError(f"Error: No atlases found matching '{args.atlas}'. Please check the atlas name and try again.")
+        elif len(df_meta) > 1:
+            raise ValueError(f"Error: Multiple atlases found matching '{args.atlas}'. Please be more specific.")
+            
+        # Read the atlas configuration from the JSON metadata
+        atlas_config = json.loads(df_meta.iloc[0]['json'])
+        
+        # Print the header with atlas information and citations
+        print_header(atlas_config['metadata'])
+        
+        # Download the enture atlas snapshot
+        atlas_local_path = hf.snapshot_download(df_meta.iloc[0]['repo'])
+                    
+    # Set up the atlas configuration the way Yue's code expects it
+    config = _setup_config(atlas_config, args, atlas_local_path, training=False)
+
+    # Create the inferencer
+    tester = HyperASHSInference(config)
+    
+    # Create links in the working directory
+    create_links, overwrite_existing = not args.no_links, not args.no_overwrite
+    for img_type, img_path in [('mprage', args.t1), ('tse', args.t2)]:        
+        dest = os.path.join(args.workdir, f'{img_type}.nii.gz')
+        copy_or_link_file(img_path, dest, create_links=create_links, force_overwrite=overwrite_existing, relative_links=False)
+            
+    # Run the segmentation 
+    print('-' * 60)   
+    print(f"Running HyperResASHS with:")
+    print(f"  Atlas:   {args.atlas}")
+    print(f"  T1:      {args.t1}")
+    print(f"  T2:      {args.t2}")
+    print(f"  Workdir: {args.workdir}")
+    print('-' * 60)
+    tester.run_inference_for_one_case(case_path=args.workdir, 
+                                      save_intermediates=args.tidy is False, 
+                                      overwrite_existing=overwrite_existing, 
+                                      create_links=create_links,
+                                      device=args.device)
+
+
+def run_training(args) -> int:
+    """Run the training pipeline."""
+    
+    create_links, overwrite_existing = not args.no_links, not args.no_overwrite
+
+    # Create a logger
+    os.makedirs(os.path.join(args.workdir, 'logs'), exist_ok=True)
+    logger = Logger(os.path.join(os.path.join(args.workdir, 'logs'), 'hyperashs_log.txt'))
+    
+    # Load the training configuration from the specified YAML file
+    try:
+        with open(os.path.join(args.config), 'r') as f:
+            atlas_config = yaml.safe_load(f)
+
+        # Print the header with atlas information and citations
+        print_header(atlas_config['metadata'])
+
+    except Exception as e:
+        print(f"Error loading atlas configuration: {e}")
+        return -1
+    
+    # Set up the atlas configuration the way Yue's code expects it
+    config = _setup_config(atlas_config, args, args.workdir, training=True)
+    
+    # Process the stage argument to determine which stages to run
+    stage_no = set()
+    if args.stage:
+        for part in args.stage.split(','):
+            if '-' in part:
+                start, end = map(int, part.split('-'))
+                stage_no.update(range(start, end + 1))
+            else:
+                stage_no.add(int(part))
+    else:
+        stage_no = {1, 2, 3, 4, 5}  # Default to running all stages if not specified
+
+    # Create the training object
+    trainer = HyperASHSTraining(config, 
+                                manifest_file=args.manifest,
+                                label_file=args.labels,
+                                xval_file=args.xval,
+                                output_dir=args.workdir,
+                                overwrite_existing=overwrite_existing,
+                                save_intermediates=args.tidy is False, 
+                                create_links=create_links)
+    
+    # Set up stages and validations for the training pipeline
+    stages = [
+        (1, trainer.preprocess, {'filter': args.filter}, None),
+        (2, trainer.train_inr, {'filter': args.filter, 'device': args.device, 'random_seed': args.inr_random_seed, 'batch_size': args.inr_batch_size}, trainer.validity_check_inr_results),
+        (3, trainer.prepare_nnunet, {}, None),
+        (4, trainer.train_nnunet, {'filter': args.filter, 'device': args.device}, trainer.validity_check_nnunet_results),
+        (5, trainer.finalize, {'full_metadata': atlas_config}, None)
+    ]
+
+    # Run the specified stages    
+    for i_stage, stage_func, stage_kwargs, validity_check in stages:
+        if i_stage in stage_no:
+            # Run the stage
+            stage_func(**stage_kwargs)
+            if max(stage_no) == i_stage:
+                return 0
+            
+        # Perform validity check at the end of this stage. This prevents moving on
+        # to the next stage if there were failures in the current stage. 
+        if validity_check is not None:
+            if validity_check() is False:
+                return 1
+  
+    return 0