ENH: Build bundle list dynamically depending on atlas version

bin/wm_append_clusters_to_anatomical_tracts.py

@@ -8,6 +8,34 @@
 import vtk
 
 import whitematteranalysis as wma
+import pandas as pd
+
+from whitematteranalysis.data.atlas.utils import (
+    ORGAtlasVersion,
+    get_local_atlas_bundle_fname,
+)
+from whitematteranalysis.anatomy.org_atlas_utils import (
+    ORGAtlasBundleFileHeading,
+    get_hemispheric_mono_bundles,
+    get_commissural_augmented_bundles,
+    add_org_atlas_prefix,
+    get_bundle_short_name,
+)
+
+
+def get_hemispheric_mono_bundle_names(_df):
+
+    hemis_mono_df = get_hemispheric_mono_bundles(_df)
+    hemis_names = hemis_mono_df[ORGAtlasBundleFileHeading.SHORT_NAME.value].values.tolist()
+    return add_org_atlas_prefix(hemis_names)
+
+
+def get_commissural_bundle_names(_df):
+
+    com_df = get_commissural_augmented_bundles(_df)
+    com_names = get_bundle_short_name(com_df, com_df[
+        ORGAtlasBundleFileHeading.ID.value].values)
+    return add_org_atlas_prefix(com_names)
 
 
 def _build_arg_parser():
@@ -24,6 +52,12 @@ def _build_arg_parser():
     parser.add_argument(
         'outputDirectory',
         help='The output directory should be a new empty directory. It will be created if needed.')
+    parser.add_argument(
+        '--version',
+        choices=ORGAtlasVersion._member_names_,
+        default=ORGAtlasVersion.V1_2.name,
+        help="Atlas version.",
+    )
 
     return parser
 
@@ -105,12 +139,11 @@ def output_appended_tract(cluster_vtp_list, outputfile):
 
         wma.io.write_polydata(pd_appended_cluster, outputfile)
 
-    hemispheric_tracts = ["T_AF", "T_CB", "T_CPC", "T_MdLF", "T_PLIC", "T_SLF-I", "T_SLF-II", "T_SLF-III", "T_EC", "T_EmC", "T_ICP", "T_ILF", "T_IOFF", "T_UF",
-                         "T_Intra-CBLM-I&P", "T_Intra-CBLM-PaT", "T_CR-F", "T_CR-P", "T_CST", "T_SF", "T_SO", "T_SP", "T_TF", "T_TO", "T_TP",
-                         "T_Sup-F", "T_Sup-FP", "T_Sup-O", "T_Sup-OT", "T_Sup-P", "T_Sup-PO", "T_Sup-PT", "T_Sup-T"]
-
-    commissural_tracts = ["T_CC1", "T_CC2", "T_CC3", "T_CC4", "T_CC5", "T_CC6", "T_CC7", "T_MCP"]
+    fname = get_local_atlas_bundle_fname(ORGAtlasVersion(args.version))
+    df = pd.read_csv(fname, sep=",")
 
+    hemispheric_tracts = get_hemispheric_mono_bundle_names(df)
+    commissural_tracts = get_commissural_bundle_names(df)
 
     print(f"<{os.path.basename(__file__)}> hemispheric tracts (left and right): ")
     tract_idx = 1

utilities/tests/test_wm_query_atlas_bundle_names.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+def test_help_option(script_runner):
+    ret = script_runner.run(["wm_query_atlas_bundles_names.py", "--help"])
+    assert ret.success

utilities/wm_query_atlas_bundles_names.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""Query the atlas bundle names.
+
+e.g. python wm_query_atlas_bundles.py \
+    /mnt/data/atlas/ORG-Atlas-1.4/ORG-800FC-100HCP/ \
+    /mnt/data/atlas/ORG-Atlas-1.4/bundle_names_default.txt \
+    DEFAULT
+
+ python wm_query_atlas_bundles.py \
+    /mnt/data/atlas/ORG-Atlas-1.4/ORG-800FC-100HCP-separated/AnatomicalTracts/ \
+    /mnt/data/atlas/ORG-Atlas-1.4/bundle_names_separated.txt \
+    SEPARATED
+"""
+
+import argparse
+from pathlib import Path
+
+from whitematteranalysis.anatomy.org_atlas_utils import (
+    AtlasAvailability,
+    query_bundle_names,
+)
+from whitematteranalysis.fileio.utils import save2txt
+
+
+def _build_arg_parser():
+
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawTextHelpFormatter
+    )
+    parser.add_argument("in_path", type=Path, help="Input path.")
+    parser.add_argument(
+        "out_fname",
+        help="Output filename (*.txt).",
+        type=Path,
+    )
+    parser.add_argument(
+        "atlas_availability",
+        choices=AtlasAvailability._member_names_,
+        help="Atlas availability.",
+    )
+
+    return parser
+
+
+def _parse_args(parser):
+
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+
+    parser = _build_arg_parser()
+    args = _parse_args(parser)
+
+    bundle_names = query_bundle_names(
+        args.in_path,
+        AtlasAvailability.__getitem__(args.atlas_availability),
+    )
+    save2txt(args.out_fname, bundle_names)
+
+
+if __name__ == "__main__":
+    main()

whitematteranalysis/anatomy/org_atlas_utils.py

@@ -0,0 +1,177 @@
+# -*- coding: utf-8 -*-
+
+import enum
+
+import pandas as pd
+
+from whitematteranalysis.fileio.utils import (
+    SlicerFileExtension,
+    VTKFileExtension,
+    find_filenames,
+)
+
+atlas_bundle_file_prefix = "T_"
+nonatlas_bundle_names = ["T_FalsePositive"]
+
+
+class Hemisphere(enum.Enum):
+    LEFT = "left"
+    RIGHT = "right"
+
+
+class AtlasAvailability(enum.Enum):
+    DEFAULT = "default"
+    SEPARATED = "separated"
+
+
+class ORGAtlasBundleType(enum.Enum):
+    ASS = "Association"
+    CER = "Cerebellar"
+    COM = "Commissural"
+    PRO = "Projection"
+    STR = "Striatal"
+    SUP = "Superficial"
+
+
+class ORGAtlasBundleFileHeading(enum.Enum):
+    ID = "id"
+    SHORT_NAME = "short_name"
+    LONG_NAME = "long_name"
+
+
+class ORGAtlasBundleRange(enum.Enum):
+    ASO = (101, 199)
+    CER = (201, 299)
+    COM = (301, 399)
+    PRO = (401, 499)
+    STR = (501, 599)
+    SUP = (601, 699)
+
+    @staticmethod
+    def get_range_low(name):
+        return ORGAtlasBundleRange(name).value[0]
+
+    @staticmethod
+    def get_range_high(name):
+        return ORGAtlasBundleRange(name).value[1]
+
+
+def get_bundles(df, bundle_id):
+    return df[df[ORGAtlasBundleFileHeading.ID.value].isin(bundle_id)]
+
+
+def get_bundle_feature(df, bundle_id, column):
+    return get_bundles(df, bundle_id)[column].values.tolist()
+
+
+def get_bundle_long_name(df, bundle_id):
+    return get_bundle_feature(df, bundle_id, ORGAtlasBundleFileHeading.LONG_NAME.value)
+
+
+def get_bundle_short_name(df, bundle_id):
+    return get_bundle_feature(df, bundle_id, ORGAtlasBundleFileHeading.SHORT_NAME.value)
+
+
+def get_bundle_type(df, bundle_type):
+    return df[
+        df[
+            ORGAtlasBundleFileHeading.ID.value].between(
+            ORGAtlasBundleRange.get_range_low(bundle_type),
+            ORGAtlasBundleRange.get_range_high(bundle_type)
+        )
+    ]
+
+
+def get_association_bundles(df):
+    return get_bundle_type(df, ORGAtlasBundleRange.ASO)
+
+
+def get_cerebellar_bundles(df):
+    return get_bundle_type(df, ORGAtlasBundleRange.CER)
+
+
+def get_commissural_bundles(df):
+    return get_bundle_type(df, ORGAtlasBundleRange.COM)
+
+
+def get_commissural_augmented_bundles(df):
+    com = get_commissural_bundles(df)
+    # Add the MCP
+    mcp = get_bundles(df, [209])
+
+    return pd.concat([com, mcp], ignore_index=True)
+
+
+def get_hemispheric_bundles(df):
+    aso = get_association_bundles(df)
+    cer = get_cerebellar_bundles(df)
+    # Remove the MCP
+    cer = cer.drop(cer[cer[ORGAtlasBundleFileHeading.ID.value] == 209].index)
+    stl = get_striatal_bundles(df)
+    pro = get_projection_bundles(df)
+    sup = get_superficial_bundles(df)
+
+    return pd.concat([aso, cer, stl, pro, sup], ignore_index=True)
+
+
+def get_hemispheric_mono_bundles(df):
+    hem_mono = get_hemispheric_bundles(df)
+
+    # Remove all right values
+    hem_mono = hem_mono.drop(hem_mono.loc[hem_mono[ORGAtlasBundleFileHeading.SHORT_NAME.value].str.contains(Hemisphere.RIGHT.value)].index)
+
+    # Drop the left substring
+    short_name_substr = "_" + Hemisphere.LEFT.value
+    column = ORGAtlasBundleFileHeading.SHORT_NAME.value
+    hem_mono[column] = hem_mono[column].str.replace(short_name_substr, "")
+    long_name_substr = " " + Hemisphere.LEFT.value
+    column = ORGAtlasBundleFileHeading.LONG_NAME.value
+    hem_mono[column] = hem_mono[column].str.replace(long_name_substr, "")
+
+    # Strip the ID as it makes no longer sense
+    hem_mono.drop(ORGAtlasBundleFileHeading.ID.value, axis=1, inplace=True)
+
+    return hem_mono
+
+
+def get_projection_bundles(df):
+    return get_bundle_type(df, ORGAtlasBundleRange.PRO)
+
+
+def get_striatal_bundles(df):
+    return get_bundle_type(df, ORGAtlasBundleRange.STR)
+
+
+def get_superficial_bundles(df):
+    return get_bundle_type(df, ORGAtlasBundleRange.SUP)
+
+
+def query_bundle_names_from_anatomical_tracts(path):
+    bundle_names = find_filenames(path, VTKFileExtension.VTP, stem=True)
+    return bundle_names
+
+
+def query_bundle_names_from_scene_files(path):
+    bundle_names = find_filenames(path, SlicerFileExtension.MRML, stem=True)
+
+    # Remove non-atlas filenames if present
+    [bundle_names.remove(elem) for elem in nonatlas_bundle_names]
+    bundle_names = [elem for elem in bundle_names if elem.startswith(atlas_bundle_file_prefix)]
+
+    return bundle_names
+
+
+def add_org_atlas_prefix(bundles):
+    return [atlas_bundle_file_prefix + elem for elem in bundles]
+
+
+def query_bundle_names(path, atlas_availability):
+
+    if atlas_availability == AtlasAvailability.DEFAULT:
+        return query_bundle_names_from_scene_files(path)
+    elif atlas_availability == AtlasAvailability.SEPARATED:
+        return query_bundle_names_from_anatomical_tracts(path)
+    else:
+        raise ValueError(
+            f"Unsupported value:\n"
+            f"Found: {atlas_availability}; Available: {AtlasAvailability.__members__}")