From f291c5ff6934543c63f7a9b2f16ee33c80052f09 Mon Sep 17 00:00:00 2001 From: guikunchen <45610535@qq.com> Date: Wed, 8 Jul 2026 16:05:47 -0400 Subject: [PATCH] ENH: Add hemisphere export mode --- pyproject.toml | 1 + src/tractcloud/README_hemisphere.md | 50 +++++ src/tractcloud/cli.py | 11 +- src/tractcloud/hemisphere.py | 323 ++++++++++++++++++++++++++++ src/tractcloud/pipeline.py | 128 +++++++---- 5 files changed, 475 insertions(+), 38 deletions(-) create mode 100644 src/tractcloud/README_hemisphere.md create mode 100644 src/tractcloud/hemisphere.py diff --git a/pyproject.toml b/pyproject.toml index 745cb34..ba778c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "numpy>=1.21", "torch>=2.3.1", "vtk>=9.0", + "whitematteranalysis @ git+https://github.com/SlicerDMRI/whitematteranalysis.git", ] [project.optional-dependencies] diff --git a/src/tractcloud/README_hemisphere.md b/src/tractcloud/README_hemisphere.md new file mode 100644 index 0000000..63af527 --- /dev/null +++ b/src/tractcloud/README_hemisphere.md @@ -0,0 +1,50 @@ +# Hemisphere Export Mode + +`tractcloud --hemisphere-atlas-dir PATH` runs TractCloud inference, keeps the +fine cluster predictions, registers the input tractography to the ORG atlas +with WMA, and writes hemisphere-split tractography outputs. + +Default CLI behavior is unchanged when `--hemisphere-atlas-dir` is omitted. +Hemisphere mode writes only Hemisphere outputs, not the default coarse +per-tract category folders. `--mrb` is not supported in this mode. + +## Requirements + +- WMA must be installed on the machine running the command (pip install git+https://github.com/SlicerDMRI/whitematteranalysis.git). +- `PATH` should point to `ORG-Atlases-1.1.1` (https://www.dropbox.com/s/beju3c0g9jqw5uj/WMA_tutorial_data.zip?dl=0) or a parent directory containing: + +```text +ORG-Atlases-1.1.1/ORG-RegAtlas-100HCP/registration_atlas.vtk +ORG-Atlases-1.1.1/ORG-800FC-100HCP/cluster_hemisphere_location.txt +``` + +## Usage + +```bash +tractcloud \ + --input subject.vtk \ + --output-dir results \ + --hemisphere-atlas-dir /path/to/ORG-Atlases-1.1.1 +``` + +## Outputs + +```text +results/ + wma_registration/ + Hemisphere/ + annotated_tractography.vtp + hemisphere_summary.json + left/ + AF_left.vtp + right/ + AF_right.vtp + commissural/ + CC1_commissural.vtp +``` + +All Hemisphere `.vtp` files include these cell arrays: + +- `FineLabel`: raw TractCloud cluster prediction. +- `CoarseID`: tract index from `tract_mapping.py`. +- `HemisphereLocation`: `0=other`, `1=left`, `2=right`, `3=commissural`. diff --git a/src/tractcloud/cli.py b/src/tractcloud/cli.py index 4d145f9..1006ac2 100644 --- a/src/tractcloud/cli.py +++ b/src/tractcloud/cli.py @@ -20,7 +20,7 @@ def main(): help="Input tractography file (.vtk or .vtp)") parser.add_argument( "--output-dir", "-o", required=True, - help="Output directory for per-tract VTP files") + help="Output directory") parser.add_argument( "--mrb", action="store_true", help="Also create a Slicer-compatible MRB file") @@ -39,12 +39,20 @@ def main(): parser.add_argument( "--data-dir", help="Override model data cache directory") + parser.add_argument( + "--hemisphere-atlas-dir", + help=( + "Run Hemisphere export mode using an ORG-Atlases-1.1.1 path. If passed, runs TractCloud inference, uses raw fine cluster_preds for hemisphere assignment, runs WMA registration, and writes only hemisphere outputs." + )) parser.add_argument( "--quiet", "-q", action="store_true", help="Suppress progress output on stdout") args = parser.parse_args() + if args.mrb and args.hemisphere_atlas_dir: + parser.error("--mrb cannot be used with --hemisphere-atlas-dir") + # Configure logging to stderr (keep stdout clean for JSON progress) logging.basicConfig( stream=sys.stderr, @@ -76,6 +84,7 @@ def main(): args.input, args.output_dir, create_mrb=args.mrb, + hemisphere_atlas_dir=args.hemisphere_atlas_dir, ) diff --git a/src/tractcloud/hemisphere.py b/src/tractcloud/hemisphere.py new file mode 100644 index 0000000..b400f37 --- /dev/null +++ b/src/tractcloud/hemisphere.py @@ -0,0 +1,323 @@ +"""Hemisphere export for TractCloud fine cluster predictions.""" + +import json +import re +import subprocess +import sys +from collections import Counter, defaultdict +from pathlib import Path + +import numpy as np +import vtk +import whitematteranalysis as wma + +from .tract_mapping import TRACT_NAMES, cluster2tract_label +from .vtk_io import extract_fibers, read_polydata, write_polydata + + +HEMI_OTHER = "other" +HEMI_LEFT = "left" +HEMI_RIGHT = "right" +HEMI_COMMISSURAL = "commissural" + +HEMI_CODE = { + HEMI_OTHER: 0, + HEMI_LEFT: 1, + HEMI_RIGHT: 2, + HEMI_COMMISSURAL: 3, +} + +WMA_H_THRESHOLD = 0.5001 +WMA_CPC_FINE_ID_FIXES = {145, 159, 557, 677, 770} + + +def export_hemisphere(input_path, input_pd, cluster_preds, output_dir, atlas_dir): + """Run WMA registration and write hemisphere-split tract outputs.""" + output_dir = Path(output_dir) + hemisphere_dir = output_dir / "Hemisphere" + registered_path = run_wma_registration(input_path, atlas_dir, output_dir) + registered_pd = read_polydata(str(registered_path)) + cluster_location_file = find_cluster_location_file(atlas_dir) + + return write_hemisphere_outputs( + input_pd=input_pd, + registered_pd=registered_pd, + cluster_preds=cluster_preds, + output_dir=hemisphere_dir, + registered_path=registered_path, + cluster_location_file=cluster_location_file, + ) + + +def run_wma_registration(input_path, atlas_dir, output_dir): + atlas_path = find_registration_atlas(atlas_dir) + registration_dir = Path(output_dir) / "wma_registration" + + cmd = [ + "wm_register_to_atlas_new.py", + "-mode", + "rigid_affine_fast", + str(input_path), + str(atlas_path), + str(registration_dir), + ] + try: + subprocess.run(cmd, check=True, stdout=sys.stderr, stderr=sys.stderr) + except FileNotFoundError as e: + raise FileNotFoundError( + "wm_register_to_atlas_new.py was not found on PATH. " + "Install WMA before running --hemisphere-atlas-dir." + ) from e + + subject_id = Path(input_path).stem + registered_path = ( + registration_dir / subject_id / "output_tractography" / f"{subject_id}_reg.vtk" + ) + if not registered_path.exists(): + raise FileNotFoundError( + f"WMA registration finished but did not write {registered_path}" + ) + return registered_path + + +def find_registration_atlas(atlas_dir): + atlas_dir = Path(atlas_dir).expanduser() + candidates = [ + atlas_dir / "ORG-Atlases-1.1.1" / "ORG-RegAtlas-100HCP" / "registration_atlas.vtk", + atlas_dir / "ORG-RegAtlas-100HCP" / "registration_atlas.vtk", + atlas_dir / "registration_atlas.vtk", + ] + for path in candidates: + if path.exists(): + return path + raise FileNotFoundError( + "Could not find registration_atlas.vtk in the provided ORG atlas path." + ) + + +def find_cluster_location_file(atlas_dir): + atlas_dir = Path(atlas_dir).expanduser() + candidates = [ + atlas_dir + / "ORG-Atlases-1.1.1" + / "ORG-800FC-100HCP" + / "cluster_hemisphere_location.txt", + atlas_dir / "ORG-800FC-100HCP" / "cluster_hemisphere_location.txt", + atlas_dir / "cluster_hemisphere_location.txt", + ] + for path in candidates: + if path.exists(): + return path + raise FileNotFoundError( + "Could not find ORG-800FC-100HCP/cluster_hemisphere_location.txt " + "in the provided ORG atlas path." + ) + + +def write_hemisphere_outputs( + input_pd, + registered_pd, + cluster_preds, + output_dir, + registered_path=None, + cluster_location_file=None, +): + input_lines = get_polyline_point_ids(input_pd) + registered_lines = get_polyline_point_ids(registered_pd) + validate_registered_tractography(input_lines, registered_lines) + + labels = np.asarray(cluster_preds).reshape(-1).astype(np.int64) + if len(labels) != len(input_lines): + raise ValueError( + f"cluster_preds length ({len(labels)}) != streamlines ({len(input_lines)})" + ) + + if cluster_location_file is None: + raise ValueError("cluster_location_file is required") + cluster_location = parse_cluster_location_file(cluster_location_file) + coarse_ids, hemi, hemi_codes = build_assignments( + registered_pd, + registered_lines, + labels, + cluster_location, + ) + + annotated = annotate_polydata(input_pd, labels, coarse_ids, hemi_codes) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + groups = defaultdict(list) + for i, hemi_name in enumerate(hemi): + tract_name = TRACT_NAMES[int(coarse_ids[i])] + if tract_name == "Other": # NOTE: discard "Other" tract + continue + groups[(hemi_name, tract_name)].append(i) + + written_files = [] + group_counts = defaultdict(dict) + for (hemi_name, tract_name), idxs in sorted(groups.items()): + group_counts[hemi_name][tract_name] = len(idxs) + out_path = output_dir / hemi_name / f"{sanitize_name(tract_name)}_{hemi_name}.vtp" + write_polydata(extract_fibers(annotated, idxs), str(out_path)) + written_files.append(str(out_path)) + + annotated_path = output_dir / "annotated_tractography.vtp" + write_polydata(annotated, str(annotated_path)) + written_files.append(str(annotated_path)) + + summary = { + "output_dir": str(output_dir), + "registered_tractography": str(registered_path) if registered_path else None, + "n_streamlines": int(len(labels)), + "hemisphere_counts": dict(Counter(hemi)), + "group_counts": {k: dict(v) for k, v in group_counts.items()}, + "written_files": written_files, + "hemisphere_code_convention": { + "0": HEMI_OTHER, + "1": HEMI_LEFT, + "2": HEMI_RIGHT, + "3": HEMI_COMMISSURAL, + }, + } + with open(output_dir / "hemisphere_summary.json", "w", encoding="utf-8") as f: + json.dump(summary, f, indent=2, sort_keys=True) + return summary + + +def parse_cluster_location_file(path): + mapping = {} + with open(path, "r", encoding="utf-8") as f: + for raw_line in f: + line = raw_line.strip() + if not line or line.lower().startswith("cluster index"): + continue + + parts = re.split(r"\s+", line) + if len(parts) < 2: + continue + + match = re.search(r"cluster_(\d+)", parts[0]) + if not match: + continue + + loc = parts[1].lower() + if loc not in {"h", "c", "ng"}: + raise ValueError(f"Unknown cluster location label {loc!r}: {line}") + mapping[int(match.group(1)) - 1] = loc + return mapping + + +def build_assignments(registered_pd, registered_lines, labels, cluster_location): + coarse_ids = np.asarray(cluster2tract_label(labels), dtype=np.int64) + hemi = np.asarray([HEMI_OTHER] * len(labels), dtype=object) + + # Use WMA FiberArray for the per-streamline geometric hemisphere test. + # This matches the core logic in wm_assess_cluster_location_by_hemisphere.py: + # points_per_fiber = 40 + # hemisphere_percent_threshold = 0.5001 for known hemispheric clusters + # RAS x/R coordinate: negative = left, positive = right + fibers = wma.fibers.FiberArray() + fibers.points_per_fiber = 40 + fibers.hemisphere_percent_threshold = WMA_H_THRESHOLD + fibers.hemispheres = True + fibers.convert_from_polydata(registered_pd) + + if fibers.number_of_fibers != len(labels): + raise ValueError( + f"WMA FiberArray fibers ({fibers.number_of_fibers}) != labels ({len(labels)})" + ) + + # Convert WMA internal codes to this script's string labels. + # WMA FiberArray convention: + # fiber_hemisphere == -1 -> left + # fiber_hemisphere == 1 -> right + # fiber_hemisphere == 0 -> commissural / ambiguous + wma_hemi = np.asarray(fibers.fiber_hemisphere, dtype=np.int64) + raw_hemi = np.asarray([HEMI_COMMISSURAL] * len(labels), dtype=object) + raw_hemi[wma_hemi == -1] = HEMI_LEFT + raw_hemi[wma_hemi == 1] = HEMI_RIGHT + raw_hemi[wma_hemi == 0] = HEMI_COMMISSURAL + + for fine_id in sorted({int(x) for x in labels}): + idxs = np.where(labels == fine_id)[0] + tract_name = TRACT_NAMES[int(coarse_ids[idxs[0]])] + + if fine_id < 0 or fine_id >= 800 or tract_name == "Other": + continue + + loc = "h" if fine_id in WMA_CPC_FINE_ID_FIXES else cluster_location.get(fine_id) + if loc is None: + raise ValueError(f"Missing cluster-location entry for fine label {fine_id}") + + if loc in {"c", "ng"}: + hemi[idxs] = HEMI_COMMISSURAL + continue + + if loc != "h": + raise ValueError(f"Unexpected cluster-location label {loc!r} for fine label {fine_id}") + + # For h clusters, use WMA FiberArray's per-streamline result first. + local = raw_hemi[idxs].copy() + + # Match WMA clusterLocationFile behavior: + # ambiguous/commissural fibers inside an h cluster are reassigned to + # the smaller side within that same cluster. + ambiguous = local == HEMI_COMMISSURAL + if np.any(ambiguous): + n_left = int(np.sum(local == HEMI_LEFT)) + n_right = int(np.sum(local == HEMI_RIGHT)) + local[ambiguous] = HEMI_LEFT if n_left <= n_right else HEMI_RIGHT + + hemi[idxs] = local + + hemi_codes = np.asarray([HEMI_CODE[h] for h in hemi], dtype=np.int64) + return coarse_ids, list(hemi), hemi_codes + + +def get_polyline_point_ids(pd): + lines = pd.GetLines() + lines.InitTraversal() + ptids = vtk.vtkIdList() + out = [] + for _ in range(pd.GetNumberOfLines()): + lines.GetNextCell(ptids) + out.append([ptids.GetId(i) for i in range(ptids.GetNumberOfIds())]) + return out + + +def validate_registered_tractography(input_lines, registered_lines): + if len(input_lines) != len(registered_lines): + raise ValueError( + "Input and registered tractography must have the same number of streamlines." + ) + + for i, (input_line, registered_line) in enumerate(zip(input_lines, registered_lines)): + if len(input_line) != len(registered_line): + raise ValueError( + "Input and registered tractography must preserve streamline order " + f"and point count. Streamline {i} differs." + ) + + +def annotate_polydata(pd, labels, coarse_ids, hemi_codes): + annotated = vtk.vtkPolyData() + annotated.DeepCopy(pd) + add_int_cell_array(annotated, "FineLabel", labels) + add_int_cell_array(annotated, "CoarseID", coarse_ids) + add_int_cell_array(annotated, "HemisphereLocation", hemi_codes) + return annotated + + +def add_int_cell_array(pd, name, values): + pd.GetCellData().RemoveArray(name) + arr = vtk.vtkIntArray() + arr.SetName(name) + arr.SetNumberOfComponents(1) + for value in values: + arr.InsertNextValue(int(value)) + pd.GetCellData().AddArray(arr) + + +def sanitize_name(name): + name = str(name).strip().replace("/", "-").replace(" ", "_") + return re.sub(r"[^A-Za-z0-9_.+-]+", "_", name) diff --git a/src/tractcloud/pipeline.py b/src/tractcloud/pipeline.py index 019c290..bab49f3 100644 --- a/src/tractcloud/pipeline.py +++ b/src/tractcloud/pipeline.py @@ -1,5 +1,6 @@ """TractCloud pipeline: end-to-end tractography parcellation.""" +import logging import time import numpy as np @@ -46,7 +47,13 @@ def __init__(self, reporter=None, device=None, batch_size=2048, import os os.environ["TRACTCLOUD_DATA_DIR"] = data_dir - def run_on_file(self, input_path, output_dir, create_mrb=False): + def run_on_file( + self, + input_path, + output_dir, + create_mrb=False, + hemisphere_atlas_dir=None, + ): """Run parcellation on a VTK/VTP file. Writes per-tract VTP files to output_dir, organized by category. @@ -55,17 +62,29 @@ def run_on_file(self, input_path, output_dir, create_mrb=False): input_path: path to input .vtk/.vtp file output_dir: directory for output files create_mrb: if True, also create an MRB file + hemisphere_atlas_dir: if set, write Hemisphere outputs only Returns: dict mapping category -> {tract_name: output_file_path} """ import os + if create_mrb and hemisphere_atlas_dir: + raise ValueError("create_mrb cannot be used with hemisphere_atlas_dir") + polydata = read_polydata(input_path) self.reporter.status( f"Loaded {polydata.GetNumberOfLines()} streamlines from " f"{os.path.basename(input_path)}") + if hemisphere_atlas_dir: + return self.run_hemisphere_on_file( + input_path, + polydata, + output_dir, + hemisphere_atlas_dir, + ) + result = self.run_on_polydata(polydata) os.makedirs(output_dir, exist_ok=True) @@ -99,6 +118,37 @@ def run_on_file(self, input_path, output_dir, create_mrb=False): return output_files + def run_hemisphere_on_file(self, input_path, polydata, output_dir, atlas_dir): + """Run inference and write Hemisphere outputs only.""" + import os + + cluster_preds, start_time = self._predict_clusters(polydata, total_steps=4) + + self.reporter.status( + "Registering tractography and writing Hemisphere outputs " + "(step 4/4)...", + step=4, + total_steps=4, + ) + from .hemisphere import export_hemisphere + + summary = export_hemisphere( + input_path=input_path, + input_pd=polydata, + cluster_preds=cluster_preds, + output_dir=output_dir, + atlas_dir=atlas_dir, + ) + self.reporter.progress(1.0, step=4) + + elapsed = time.time() - start_time + self.reporter.result( + len(summary["written_files"]), + os.path.join(output_dir, "Hemisphere"), + total_time=elapsed, + ) + return summary + def run_on_polydata(self, polydata): """Run parcellation on a vtkPolyData. @@ -108,10 +158,47 @@ def run_on_polydata(self, polydata): Returns: dict: {category_name: {tract_name: vtkPolyData}} """ - start_time = time.time() total_steps = 4 + cluster_preds, start_time = self._predict_clusters( + polydata, + total_steps=total_steps, + ) + + # Step 4: Output + self.reporter.status( + f"Extracting tract bundles (step 4/{total_steps})...", + step=4, total_steps=total_steps) + tract_labels = np.array(cluster2tract_label(cluster_preds)) + + result = {} + total_tracts = sum( + len(tracts) for cat, tracts in TRACT_CATEGORIES.items() + if cat != "Other" or self.include_other) + tracts_done = 0 + + for category, tract_names_in_cat in TRACT_CATEGORIES.items(): + if category == "Other" and not self.include_other: + continue + cat_tracts = {} + for tract_name in tract_names_in_cat: + tract_idx = TRACT_NAMES.index(tract_name) + fiber_idx = np.where(tract_labels == tract_idx)[0] + if len(fiber_idx) == 0: + continue + cat_tracts[tract_name] = extract_fibers(polydata, fiber_idx) + tracts_done += 1 + self.reporter.progress(tracts_done / total_tracts, step=4) + if cat_tracts: + result[category] = cat_tracts + + elapsed = time.time() - start_time + total_created = sum(len(t) for t in result.values()) + self.reporter.result(total_created, "", total_time=elapsed) + return result + + def _predict_clusters(self, polydata, total_steps): + start_time = time.time() - # Report device selection if self.device.type == "cuda": torch.backends.cudnn.enabled = False self.reporter.status( @@ -121,7 +208,6 @@ def run_on_polydata(self, polydata): "No GPU detected, using CPU. " "Inference will be significantly slower.") - # Step 1: Extract features num_fibers = polydata.GetNumberOfLines() self.reporter.status( f"Extracting features from {num_fibers} streamlines " @@ -139,7 +225,6 @@ def run_on_polydata(self, polydata): feat = extract_ras_features(polydata, num_points=self.num_points) self.reporter.progress(1.0, step=1) - # Step 2: KNN self.reporter.status( "Re-centering and computing neighbor features " f"(step 2/{total_steps}, slowest step)...", @@ -152,7 +237,6 @@ def run_on_polydata(self, polydata): data_loader = torch.utils.data.DataLoader( dataset, batch_size=self.batch_size, shuffle=False) - # Step 3: Inference (with CPU fallback for incompatible GPUs) self.reporter.status( f"Running model inference on {self.device} " f"(step 3/{total_steps})...", @@ -181,37 +265,7 @@ def run_on_polydata(self, polydata): else: raise - # Step 4: Output - self.reporter.status( - f"Extracting tract bundles (step 4/{total_steps})...", - step=4, total_steps=total_steps) - tract_labels = np.array(cluster2tract_label(cluster_preds)) - - result = {} - total_tracts = sum( - len(tracts) for cat, tracts in TRACT_CATEGORIES.items() - if cat != "Other" or self.include_other) - tracts_done = 0 - - for category, tract_names_in_cat in TRACT_CATEGORIES.items(): - if category == "Other" and not self.include_other: - continue - cat_tracts = {} - for tract_name in tract_names_in_cat: - tract_idx = TRACT_NAMES.index(tract_name) - fiber_idx = np.where(tract_labels == tract_idx)[0] - if len(fiber_idx) == 0: - continue - cat_tracts[tract_name] = extract_fibers(polydata, fiber_idx) - tracts_done += 1 - self.reporter.progress(tracts_done / total_tracts, step=4) - if cat_tracts: - result[category] = cat_tracts - - elapsed = time.time() - start_time - total_created = sum(len(t) for t in result.values()) - self.reporter.result(total_created, "", total_time=elapsed) - return result + return cluster_preds, start_time def run_on_features(self, feat_ras): """Run on pre-extracted numpy features.