diff --git a/README.md b/README.md index 2d33e7f..932abce 100644 --- a/README.md +++ b/README.md @@ -297,7 +297,6 @@ from physiomotion4d import RegisterImagesICON, WorkflowConvertImageToUSD # Initialize processor processor = WorkflowConvertImageToUSD( input_filenames=["path/to/cardiac_4d_ct.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", registration_method=RegisterImagesICON(), # or RegisterImagesANTS() @@ -334,20 +333,21 @@ registered_mesh = result["registered_template_model_surface"] ### Custom Segmentation ```python -from physiomotion4d import SegmentChestTotalSegmentator +from physiomotion4d import SegmentChestTotalSegmentatorWithContrast import itk -# Initialize TotalSegmentator segmentation -segmenter = SegmentChestTotalSegmentator() +# Use SegmentChestTotalSegmentator instead for non-contrast studies +segmenter = SegmentChestTotalSegmentatorWithContrast() # Load and segment image image = itk.imread("chest_ct.nrrd") -masks = segmenter.segment(image, contrast_enhanced_study=True) +masks = segmenter.segment(image) # Result always contains "labelmap" plus one entry per anatomy group the # segmenter registered (heart, lung, bone, major_vessels, soft_tissue, -# contrast, other for SegmentChestTotalSegmentator). The exact key set is -# segmenter-specific; check membership when targeting multiple segmenters. +# contrast, other for SegmentChestTotalSegmentatorWithContrast). The exact +# key set is segmenter-specific; check membership when targeting multiple +# segmenters. labelmap = masks["labelmap"] heart_mask = masks["heart"] if "lung" in masks: diff --git a/docs/api/segmentation/base.rst b/docs/api/segmentation/base.rst index 33d223a..66309e0 100644 --- a/docs/api/segmentation/base.rst +++ b/docs/api/segmentation/base.rst @@ -31,7 +31,7 @@ Concrete segmenters accept an ITK image and return a dictionary of ITK images: image = itk.imread("chest_ct.nrrd") segmenter = SegmentChestTotalSegmentator() - masks = segmenter.segment(image, contrast_enhanced_study=True) + masks = segmenter.segment(image) labelmap = masks["labelmap"] if "heart" in masks: diff --git a/docs/api/segmentation/index.rst b/docs/api/segmentation/index.rst index bde3c2c..6c74902 100644 --- a/docs/api/segmentation/index.rst +++ b/docs/api/segmentation/index.rst @@ -46,7 +46,7 @@ Basic Segmentation from physiomotion4d import SegmentChestTotalSegmentator segmenter = SegmentChestTotalSegmentator() - result = segmenter.segment(ct_image, contrast_enhanced_study=False) + result = segmenter.segment(ct_image) labelmap = result['labelmap'] Module Documentation diff --git a/docs/api/segmentation/simpleware.rst b/docs/api/segmentation/simpleware.rst index 887258a..0ee5889 100644 --- a/docs/api/segmentation/simpleware.rst +++ b/docs/api/segmentation/simpleware.rst @@ -28,7 +28,7 @@ Basic Usage image = itk.imread("chest_ct.nrrd") segmenter = SegmentHeartSimpleware() - masks = segmenter.segment(image, contrast_enhanced_study=True) + masks = segmenter.segment(image) heart = masks["heart"] vessels = masks["major_vessels"] diff --git a/docs/api/segmentation/totalsegmentator.rst b/docs/api/segmentation/totalsegmentator.rst index 4f38e21..66db952 100644 --- a/docs/api/segmentation/totalsegmentator.rst +++ b/docs/api/segmentation/totalsegmentator.rst @@ -28,7 +28,7 @@ Basic Usage image = itk.imread("chest_ct.nrrd") segmenter = SegmentChestTotalSegmentator() - masks = segmenter.segment(image, contrast_enhanced_study=True) + masks = segmenter.segment(image) heart = masks["heart"] lungs = masks["lung"] @@ -53,15 +53,30 @@ keys: * ``bone`` * ``soft_tissue`` * ``other`` -* ``contrast`` The dictionary should be accessed by key. Do not unpack it positionally. The exact key set is determined by the segmenter's :class:`AnatomyTaxonomy` and may differ from other segmenters (see :doc:`base`). For -:class:`SegmentChestTotalSegmentator` specifically, all seven groups plus +:class:`SegmentChestTotalSegmentator` specifically, all six groups plus ``labelmap`` are always present; downstream code that targets a different segmenter should check membership. +For contrast-enhanced studies, use +:class:`SegmentChestTotalSegmentatorWithContrast` instead of +:class:`SegmentChestTotalSegmentator`. It adds a ``contrast`` key to the +returned dictionary and exposes a ``contrast_threshold`` attribute +(default 500) that can be overridden before calling ``segment()``: + +.. code-block:: python + + from physiomotion4d import SegmentChestTotalSegmentatorWithContrast + + segmenter = SegmentChestTotalSegmentatorWithContrast() + segmenter.contrast_threshold = 600 # optional override + + masks = segmenter.segment(image) + contrast = masks["contrast"] + Operational Notes ================= diff --git a/docs/api/workflows.rst b/docs/api/workflows.rst index cd8fd79..3ad644f 100644 --- a/docs/api/workflows.rst +++ b/docs/api/workflows.rst @@ -50,16 +50,15 @@ Convert Image to USD from physiomotion4d import ( RegisterImagesICON, - SegmentChestTotalSegmentator, + SegmentChestTotalSegmentatorWithContrast, WorkflowConvertImageToUSD, ) workflow = WorkflowConvertImageToUSD( input_filenames=["cardiac_4d.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="patient_001", - segmentation_method=SegmentChestTotalSegmentator(), + segmentation_method=SegmentChestTotalSegmentatorWithContrast(), registration_method=RegisterImagesICON(), ) @@ -77,15 +76,17 @@ Image to VTK import itk - from physiomotion4d import SegmentChestTotalSegmentator, WorkflowConvertImageToVTK + from physiomotion4d import ( + SegmentChestTotalSegmentatorWithContrast, + WorkflowConvertImageToVTK, + ) image = itk.imread("chest_ct.nii.gz") workflow = WorkflowConvertImageToVTK( - segmentation_method=SegmentChestTotalSegmentator() + segmentation_method=SegmentChestTotalSegmentatorWithContrast() ) result = workflow.run_workflow( input_image=image, - contrast_enhanced_study=True, anatomy_groups=["heart", "major_vessels"], ) diff --git a/docs/cli_scripts/byod_tutorials.rst b/docs/cli_scripts/byod_tutorials.rst index 87b560b..1823b0c 100644 --- a/docs/cli_scripts/byod_tutorials.rst +++ b/docs/cli_scripts/byod_tutorials.rst @@ -92,7 +92,7 @@ scene. workflow = pm4d.WorkflowConvertImageToUSD( input_filenames=["patient_dicom_dir"], - contrast_enhanced=False, + segmentation_method=pm4d.SegmentChestTotalSegmentator(), output_directory="./results", project_name="patient_heart", ) @@ -139,7 +139,7 @@ selects its default reference frame internally. workflow = pm4d.WorkflowConvertImageToUSD( input_filenames=["phase_000.mha", "phase_001.mha", "phase_002.mha"], - contrast_enhanced=False, + segmentation_method=pm4d.SegmentChestTotalSegmentator(), output_directory="./results", project_name="heart_animated", times_per_second=30.0, diff --git a/docs/developer/segmentation.rst b/docs/developer/segmentation.rst index 9bf628e..98073ec 100644 --- a/docs/developer/segmentation.rst +++ b/docs/developer/segmentation.rst @@ -16,7 +16,7 @@ Current Segmentation Contract image = itk.imread("chest_ct.nrrd") segmenter = SegmentChestTotalSegmentator() - masks = segmenter.segment(image, contrast_enhanced_study=True) + masks = segmenter.segment(image) labelmap = masks["labelmap"] if "heart" in masks: diff --git a/docs/developer/workflows.rst b/docs/developer/workflows.rst index 38e66a1..e793dc0 100644 --- a/docs/developer/workflows.rst +++ b/docs/developer/workflows.rst @@ -36,7 +36,6 @@ Workflow Example workflow = WorkflowConvertImageToUSD( input_filenames=["cardiac_4d.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="patient_001", registration_method=RegisterImagesICON(), diff --git a/docs/examples.rst b/docs/examples.rst index 2bd85d0..5d155d7 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -32,7 +32,6 @@ Complete end-to-end cardiac CT processing: # Initialize workflow workflow = WorkflowConvertImageToUSD( input_filenames=["cardiac_4d.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="patient_001", registration_method=RegisterImagesICON(), @@ -83,15 +82,15 @@ Quick segmentation with TotalSegmentator: .. code-block:: python - from physiomotion4d import SegmentChestTotalSegmentator + from physiomotion4d import SegmentChestTotalSegmentatorWithContrast import itk # Load image image = itk.imread("chest_ct.nrrd") - # Segment - segmenter = SegmentChestTotalSegmentator() - masks = segmenter.segment(image, contrast_enhanced_study=True) + # Segment (use SegmentChestTotalSegmentator instead for non-contrast studies) + segmenter = SegmentChestTotalSegmentatorWithContrast() + masks = segmenter.segment(image) # Extract masks by anatomy group heart = masks["heart"] @@ -409,7 +408,6 @@ Batch process multiple datasets: workflow = WorkflowConvertImageToUSD( input_filenames=[input_file], - contrast_enhanced=True, output_directory=f"results/{patient_id}", project_name=patient_id, registration_method=RegisterImagesICON(), @@ -428,15 +426,15 @@ Segment multiple images in parallel: .. code-block:: python - from physiomotion4d import SegmentChestTotalSegmentator + from physiomotion4d import SegmentChestTotalSegmentatorWithContrast import itk import glob from concurrent.futures import ProcessPoolExecutor def segment_image(filename): - segmenter = SegmentChestTotalSegmentator() + segmenter = SegmentChestTotalSegmentatorWithContrast() image = itk.imread(filename) - result = segmenter.segment(image, contrast_enhanced_study=True) + result = segmenter.segment(image) # Save heart mask output_name = filename.replace('.nrrd', '_heart.nrrd') @@ -479,7 +477,6 @@ Run the supported end-to-end workflow API: workflow = WorkflowConvertImageToUSD( input_filenames=["cardiac_4d.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", registration_method=RegisterImagesICON(), @@ -497,7 +494,7 @@ Mix and match different components: .. code-block:: python from physiomotion4d import ( - SegmentChestTotalSegmentator, + SegmentChestTotalSegmentatorWithContrast, RegisterImagesICON, TransformTools, ConvertVTKToUSD, @@ -510,8 +507,8 @@ Mix and match different components: frames = [itk.imread(f"frame_{i:03d}.mha") for i in range(10)] # Segment reference - segmenter = SegmentChestTotalSegmentator() - result = segmenter.segment(reference, contrast_enhanced_study=True) + segmenter = SegmentChestTotalSegmentatorWithContrast() + result = segmenter.segment(reference) heart_mask = result['heart'] # Extract reference contour diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 8e637df..09f3d75 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -117,7 +117,6 @@ For more control, use the Python API: processor = WorkflowConvertImageToUSD( input_filenames=["path/to/cardiac_4d_ct.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", registration_method=RegisterImagesICON(), @@ -152,7 +151,6 @@ For more control over individual steps: # Initialize workflow workflow = WorkflowConvertImageToUSD( input_filenames=["cardiac_4d.nrrd"], - contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", registration_method=RegisterImagesICON(), @@ -170,15 +168,15 @@ If you only need segmentation: .. code-block:: python - from physiomotion4d import SegmentChestTotalSegmentator + from physiomotion4d import SegmentChestTotalSegmentatorWithContrast import itk - # Initialize segmenter - segmenter = SegmentChestTotalSegmentator() + # Initialize segmenter (use SegmentChestTotalSegmentator for non-contrast studies) + segmenter = SegmentChestTotalSegmentatorWithContrast() # Load and segment image image = itk.imread("chest_ct.nrrd") - masks = segmenter.segment(image, contrast_enhanced_study=True) + masks = segmenter.segment(image) # Extract individual anatomy masks by key heart_mask = masks["heart"] diff --git a/docs/troubleshooting.rst b/docs/troubleshooting.rst index fc9e8a0..48c004d 100644 --- a/docs/troubleshooting.rst +++ b/docs/troubleshooting.rst @@ -69,13 +69,17 @@ Poor Segmentation Quality **Solutions**: -1. Check if image is contrast-enhanced: +1. Check if image is contrast-enhanced. Use + :class:`SegmentChestTotalSegmentatorWithContrast` instead of + :class:`SegmentChestTotalSegmentator` for contrast-enhanced studies: .. code-block:: python + from physiomotion4d import SegmentChestTotalSegmentatorWithContrast + workflow = WorkflowConvertImageToUSD( ..., - contrast_enhanced=True # or False + segmentation_method=SegmentChestTotalSegmentatorWithContrast(), ) 2. Preprocess intensity, spacing, and field of view before invoking the workflow. diff --git a/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/3-eval_icon.py b/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/3-eval_icon.py index 95858e3..42da215 100644 --- a/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/3-eval_icon.py +++ b/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/3-eval_icon.py @@ -351,7 +351,7 @@ def _build_registrar( compression=True, ) - seg_result = segmenter.segment(warped_ref, contrast_enhanced_study=False) + seg_result = segmenter.segment(warped_ref) warped_ref_labelmap = seg_result["labelmap"] warped_ref_landmarks = segmenter.get_landmarks() diff --git a/experiments/Heart-GatedCT_To_USD/1-register_images.py b/experiments/Heart-GatedCT_To_USD/1-register_images.py index f392168..153d175 100644 --- a/experiments/Heart-GatedCT_To_USD/1-register_images.py +++ b/experiments/Heart-GatedCT_To_USD/1-register_images.py @@ -5,7 +5,9 @@ import itk from physiomotion4d.register_images_ants import RegisterImagesANTS -from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator +from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from physiomotion4d.test_tools import TestTools from physiomotion4d.transform_tools import TransformTools @@ -32,10 +34,8 @@ fixed_image = itk.imread(str(fixed_image_filename)) # %% - seg = SegmentChestTotalSegmentator() - seg.contrast_threshold = 500 + seg = SegmentChestTotalSegmentatorWithContrast() seg.fast_mode = test_mode - seg.set_contrast_enhanced_study(True) result = seg.segment(fixed_image) # %% labelmap_mask = result["labelmap"] diff --git a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py index 3b127ab..a5dbc10 100644 --- a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py +++ b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py @@ -7,7 +7,9 @@ import pyvista as pv from physiomotion4d.contour_tools import ContourTools -from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator +from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from physiomotion4d.test_tools import TestTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows @@ -68,10 +70,8 @@ ) outname = "slice_max" - seg = SegmentChestTotalSegmentator() - seg.contrast_threshold = 500 + seg = SegmentChestTotalSegmentatorWithContrast() seg.fast_mode = TestTools.running_as_test() - seg.set_contrast_enhanced_study(True) if re_run_image_segmentation: result = seg.segment(max_image) labelmap_image = result["labelmap"] diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py index 4376127..3ac6bcd 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py @@ -59,10 +59,7 @@ # %% if False: segmentator = SegmentChestTotalSegmentator() - segmentator.contrast_threshold = 500 - patient_segmentation_data = segmentator.segment( - patient_image, contrast_enhanced_study=False - ) + patient_segmentation_data = segmentator.segment(patient_image) labelmap = patient_segmentation_data["labelmap"] lung_mask = patient_segmentation_data["lung"] heart_mask = patient_segmentation_data["heart"] @@ -70,14 +67,12 @@ bone_mask = patient_segmentation_data["bone"] soft_tissue_mask = patient_segmentation_data["soft_tissue"] other_mask = patient_segmentation_data["other"] - contrast_mask = patient_segmentation_data["contrast"] itk.imwrite( labelmap, str(output_dir / "patient_labelmap.mha"), compression=True ) heart_arr = itk.GetArrayFromImage(heart_mask) - # contrast_arr = itk.GetArrayFromImage(contrast_mask) mask_arr = (heart_arr > 0).astype( np.uint8 ) # ((heart_arr + contrast_arr) > 0).astype(np.uint8) diff --git a/experiments/Heart-VTKSeries_To_USD/1-heart_vtkseries_to_usd.py b/experiments/Heart-VTKSeries_To_USD/1-heart_vtkseries_to_usd.py index 4fc042f..147a888 100644 --- a/experiments/Heart-VTKSeries_To_USD/1-heart_vtkseries_to_usd.py +++ b/experiments/Heart-VTKSeries_To_USD/1-heart_vtkseries_to_usd.py @@ -21,18 +21,17 @@ import itk from physiomotion4d.contour_tools import ContourTools - from physiomotion4d.segment_chest_total_segmentator import ( - SegmentChestTotalSegmentator, + from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, ) input_images = sorted(glob.glob(os.path.join(_DATA_DIR, "slice_*.mha"))) - seg = SegmentChestTotalSegmentator() - seg.contrast_threshold = 500 + seg = SegmentChestTotalSegmentatorWithContrast() con = ContourTools() for i, img_path in enumerate(input_images): print(f"Segmenting {img_path}...") img = itk.imread(img_path) - result = seg.segment(img, contrast_enhanced_study=True) + result = seg.segment(img) labelmap_mask = result["labelmap"] img_con = con.extract_contours(labelmap_mask) img_con.save(os.path.join(_DATA_DIR, f"slice_{i:03d}.vtp")) diff --git a/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py b/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py index 54fa954..f2d5298 100644 --- a/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py +++ b/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py @@ -42,7 +42,7 @@ # %% img = itk.imread(os.path.join(_RESULTS_DIR, "Experiment_reg.mha")) tot_seg = SegmentChestTotalSegmentator() - seg_results = tot_seg.segment(img, contrast_enhanced_study=False) + seg_results = tot_seg.segment(img) itk.imwrite( seg_results["labelmap"], os.path.join(_RESULTS_DIR, "Experiment_totseg.mha"), diff --git a/src/physiomotion4d/__init__.py b/src/physiomotion4d/__init__.py index cf90752..73cdd5e 100644 --- a/src/physiomotion4d/__init__.py +++ b/src/physiomotion4d/__init__.py @@ -64,6 +64,9 @@ # Segmentation classes from .segment_anatomy_base import SegmentAnatomyBase from .segment_chest_total_segmentator import SegmentChestTotalSegmentator +from .segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from .segment_heart_simpleware import SegmentHeartSimpleware from .segment_heart_simpleware_trimmed_branches import ( SegmentHeartSimplewareTrimmedBranches, @@ -96,6 +99,7 @@ # Segmentation classes "SegmentAnatomyBase", "SegmentChestTotalSegmentator", + "SegmentChestTotalSegmentatorWithContrast", "SegmentHeartSimpleware", "SegmentHeartSimplewareTrimmedBranches", # Registration classes diff --git a/src/physiomotion4d/cli/_method_factories.py b/src/physiomotion4d/cli/_method_factories.py index 59c13f8..eb4191d 100644 --- a/src/physiomotion4d/cli/_method_factories.py +++ b/src/physiomotion4d/cli/_method_factories.py @@ -12,6 +12,7 @@ RegisterImagesICON, SegmentAnatomyBase, SegmentChestTotalSegmentator, + SegmentChestTotalSegmentatorWithContrast, SegmentHeartSimpleware, SegmentHeartSimplewareTrimmedBranches, ) @@ -27,20 +28,31 @@ REGISTRATION_METHODS: tuple[str, ...] = ("Greedy", "ICON", "Greedy_ICON") -def build_segmentation_method(name: str) -> SegmentAnatomyBase: +def build_segmentation_method(name: str, contrast: bool = False) -> SegmentAnatomyBase: """Build a SegmentAnatomyBase instance for a CLI --segmentation-method choice. Args: name: One of SEGMENTATION_METHODS. + contrast: If True, build the contrast-enhanced variant instead of the + plain backend. Only supported for "ChestTotalSegmentator". Returns: A new, unconfigured segmentation backend instance. Raises: - ValueError: If name is not one of SEGMENTATION_METHODS. + ValueError: If name is not one of SEGMENTATION_METHODS, or if + contrast=True is requested for a backend that has no + contrast-enhanced variant. """ if name == "ChestTotalSegmentator": + if contrast: + return SegmentChestTotalSegmentatorWithContrast() return SegmentChestTotalSegmentator() + if contrast: + raise ValueError( + f"contrast=True is not supported for segmentation method: {name}. " + "Only ChestTotalSegmentator has a contrast-enhanced variant." + ) if name == "HeartSimpleware": return SegmentHeartSimpleware() if name == "HeartSimplewareTrimmedBranches": diff --git a/src/physiomotion4d/cli/convert_image_to_usd.py b/src/physiomotion4d/cli/convert_image_to_usd.py index 8a46f44..0ac5f15 100644 --- a/src/physiomotion4d/cli/convert_image_to_usd.py +++ b/src/physiomotion4d/cli/convert_image_to_usd.py @@ -9,14 +9,12 @@ import argparse import os import sys -from typing import cast import itk from ..convert_image_4d_to_3d import ConvertImage4DTo3D from ..register_images_greedy import RegisterImagesGreedy from ..register_images_icon import RegisterImagesICON -from ..segment_chest_total_segmentator import SegmentChestTotalSegmentator from ._method_factories import build_registration_method, build_segmentation_method @@ -121,13 +119,9 @@ def main() -> int: try: from .. import WorkflowConvertImageToUSD - segmentation_method = build_segmentation_method(args.segmentation_method) - if args.segmentation_method == "ChestTotalSegmentator": - segmentation_method_tot = cast( - SegmentChestTotalSegmentator, segmentation_method - ) - segmentation_method_tot.set_contrast_enhanced_study(args.contrast) - segmentation_method_tot.contrast_threshold = 500 + segmentation_method = build_segmentation_method( + args.segmentation_method, contrast=args.contrast + ) registration_method = build_registration_method(args.registration_method) if ( args.registration_iterations is not None @@ -145,8 +139,6 @@ def main() -> int: registration_method.set_number_of_iterations( args.registration_iterations ) - if isinstance(registration_method, RegisterImagesICON): - registration_method.set_mass_preservation(not args.contrast) if len(args.input_files) == 1: convert_image_4d_to_3d = ConvertImage4DTo3D() diff --git a/src/physiomotion4d/cli/convert_image_to_vtk.py b/src/physiomotion4d/cli/convert_image_to_vtk.py index f300f43..df64822 100644 --- a/src/physiomotion4d/cli/convert_image_to_vtk.py +++ b/src/physiomotion4d/cli/convert_image_to_vtk.py @@ -159,11 +159,12 @@ def main() -> int: from .. import WorkflowConvertImageToVTK workflow = WorkflowConvertImageToVTK( - segmentation_method=build_segmentation_method(args.segmentation_method), + segmentation_method=build_segmentation_method( + args.segmentation_method, contrast=args.contrast + ), ) result = workflow.run_workflow( input_image=input_image, - contrast_enhanced_study=args.contrast, anatomy_groups=args.anatomy_groups, ) except (ValueError, RuntimeError, OSError) as exc: diff --git a/src/physiomotion4d/segment_chest_total_segmentator.py b/src/physiomotion4d/segment_chest_total_segmentator.py index cbc48ac..cd0bc75 100644 --- a/src/physiomotion4d/segment_chest_total_segmentator.py +++ b/src/physiomotion4d/segment_chest_total_segmentator.py @@ -9,13 +9,11 @@ import logging import os import tempfile -from typing import Optional import itk import nibabel as nib import numpy as np -from .image_tools import ImageTools from .segment_anatomy_base import SegmentAnatomyBase @@ -38,15 +36,17 @@ class SegmentChestTotalSegmentator(SegmentAnatomyBase): consumers (``ConvertVTKToUSD``, ``USDAnatomyTools``) see a single, consistent group→organ mapping. + For contrast-enhanced studies (CT with contrast-enhanced blood in the + heart/vessels), use :class:`SegmentChestTotalSegmentatorWithContrast` + instead, which subclasses this class and adds a connected-component + pass to label contrast-enhanced blood under a ``"contrast"`` taxonomy + group. + Attributes: target_spacing (float): Target spacing set to 1.5mm for TotalSegmentator. - contrast_threshold (int): Lower intensity threshold used to detect - contrast-enhanced blood when contrast-enhanced-study detection - is enabled via :meth:`set_contrast_enhanced_study`. Example: >>> segmenter = SegmentChestTotalSegmentator() - >>> segmenter.set_contrast_enhanced_study(True) >>> result = segmenter.segment(ct_image) >>> labelmap = result['labelmap'] >>> heart_mask = result['heart'] @@ -67,9 +67,6 @@ def __init__(self, log_level: int | str = logging.INFO): self.target_spacing = 1.5 - self.contrast_enhanced_study: bool = False - self.contrast_threshold: int = 700 - # TotalSegmentator class indices, grouped by anatomy. for group_name, organs in ( ( @@ -208,16 +205,24 @@ def __init__(self, log_level: int | str = logging.INFO): 133: "soft_tissue", }, ), - ( - "contrast", - {135: "contrast"}, - ), ): for label_id, organ_name in organs.items(): self.taxonomy.add_organ(group_name, label_id, organ_name) + self._add_extra_taxonomy_groups() self._finalize_other_group() + def _add_extra_taxonomy_groups(self) -> None: + """Hook for subclasses to add taxonomy groups before finalization. + + Called at the end of :meth:`__init__`, before + :meth:`SegmentAnatomyBase._finalize_other_group` claims unclaimed ids + into the ``other`` group. Subclasses (e.g. + :class:`SegmentChestTotalSegmentatorWithContrast`) override this to + register additional groups without duplicating the base class's + organ mapping. + """ + def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: """ Run TotalSegmentator on the preprocessed image and return result. @@ -323,196 +328,3 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: labelmap_image.CopyInformation(preprocessed_image) return labelmap_image - - def set_contrast_enhanced_study(self, contrast_enhanced_study: bool) -> None: - """Enable or disable contrast-enhanced-study detection. - - When enabled, :meth:`segment` runs an additional connected-component - pass (see :meth:`segment_contrast_agent`) to identify contrast-enhanced - blood vessels and cardiac chambers, labeling them under the - ``"contrast"`` taxonomy group. - - Args: - contrast_enhanced_study (bool): Whether the study uses contrast - enhancement. - - Example: - >>> segmenter.set_contrast_enhanced_study(True) - """ - self.contrast_enhanced_study = contrast_enhanced_study - - def postprocess_after_labelmap( - self, input_image: itk.image, labelmap_image: itk.image - ) -> itk.image: - """Run contrast-enhanced-study detection when enabled. - - Overrides :meth:`SegmentAnatomyBase.postprocess_after_labelmap`. - - Args: - input_image (itk.image): The original, unpreprocessed input image - labelmap_image (itk.image): The postprocessed segmentation labelmap - - Returns: - itk.image: The labelmap, with contrast-enhanced regions labeled - if :attr:`contrast_enhanced_study` is True; unchanged otherwise - """ - if self.contrast_enhanced_study: - return self.segment_contrast_agent(input_image, labelmap_image) - return labelmap_image - - def segment_connected_component( - self, - preprocessed_image: itk.image, - labelmap_image: itk.image, - lower_threshold: int, - upper_threshold: int, - labelmap_ids: Optional[list[int]] = None, - mask_id: int = 0, - use_mid_slice: bool = True, - hole_fill: int = 2, - ) -> itk.image: - """ - Segment connected components based on intensity thresholding. - - Identifies connected regions within intensity thresholds and existing - anatomical masks, then selects the largest component. This is useful - for segmenting structures like contrast-enhanced blood or specific - tissue types. - - Args: - preprocessed_image (itk.image): The preprocessed input image - labelmap_image (itk.image): Existing labelmap to constrain search - lower_threshold (int): Lower intensity threshold - upper_threshold (int): Upper intensity threshold - labelmap_ids (Optional[list[int]]): List of label IDs to search within. - If None, searches within all existing labels - mask_id (int): ID to assign to the segmented component - use_mid_slice (bool): If True, find largest component in middle - slice only; if False, use entire 3D volume - hole_fill (int): Number of pixels to dilate/erode for hole filling - - Returns: - itk.image: Updated labelmap with new component labeled as mask_id - - Example: - >>> # Segment contrast-enhanced blood - >>> updated_labels = segmenter.segment_connected_component( - ... preprocessed_image, labels, 700, 4000, mask_id=135 - ... ) - """ - thresh_image = itk.binary_threshold_image_filter( - Input=preprocessed_image, - LowerThreshold=lower_threshold, - UpperThreshold=upper_threshold, - InsideValue=1, - OutsideValue=0, - ) - thresh_arr = itk.GetArrayFromImage(thresh_image).astype(np.int16) - thresh_image = itk.GetImageFromArray(thresh_arr) - thresh_image.CopyInformation(preprocessed_image) - - label_arr = itk.GetArrayFromImage(labelmap_image) - if labelmap_ids is None: - labelmap_ids = list(self.taxonomy.all_labels().keys()) - label_arr = np.isin(label_arr, labelmap_ids) - label_image = itk.GetImageFromArray(label_arr.astype(np.int16)) - label_image.CopyInformation(labelmap_image) - - connected_component_image = itk.connected_component_image_filter( - Input=thresh_image, - MaskImage=label_image, - ) - - connected_component_arr = itk.GetArrayFromImage(connected_component_image) - if use_mid_slice: - mid_slice = ( - connected_component_image.GetLargestPossibleRegion().GetSize()[2] // 2 - ) - tmp_connected_component_arr = connected_component_arr[mid_slice, :, :] - ids = np.unique(tmp_connected_component_arr) - if len(ids[ids != 0]) > 0: - connected_component_arr = tmp_connected_component_arr - - ids = np.unique(connected_component_arr) - ids = ids[ids != 0] - if ids.size == 0: - self.log_debug( - "segment_connected_component: no connected components found " - "in threshold [%d, %d]; returning labelmap unchanged", - lower_threshold, - upper_threshold, - ) - return labelmap_image - component_sums = [np.sum(connected_component_arr == id) for id in ids] - largest_id = ids[np.argmax(component_sums)] - connected_component_image = itk.binary_threshold_image_filter( - Input=connected_component_image, - LowerThreshold=int(largest_id), - UpperThreshold=int(largest_id), - InsideValue=1, - OutsideValue=0, - ) - image_tools = ImageTools() - connected_component_image = image_tools.binary_dilate_image( - connected_component_image, hole_fill, 1, 0 - ) - connected_component_image = image_tools.binary_erode_image( - connected_component_image, hole_fill, 1, 0 - ) - - labelmap_arr = itk.GetArrayFromImage(labelmap_image) - connected_component_arr = itk.GetArrayFromImage(connected_component_image) - connected_component_mask = connected_component_arr > 0 - mask = label_arr & connected_component_mask - labelmap_arr = np.where(mask, mask_id, labelmap_arr) - results_image = itk.GetImageFromArray(labelmap_arr.astype(np.uint8)) - results_image.CopyInformation(preprocessed_image) - - return results_image - - def segment_contrast_agent( - self, preprocessed_image: itk.image, labelmap_image: itk.image - ) -> itk.image: - """ - Include contrast-enhanced blood in the labelmap. - - Segments high-intensity regions corresponding to contrast-enhanced - blood vessels and cardiac chambers. Uses connected component analysis - focused on the middle slice where the heart is typically located. - - Args: - preprocessed_image (itk.image): The preprocessed CT image - labelmap_image (itk.image): Existing segmentation labelmap - - Returns: - itk.image: Updated labelmap with contrast-enhanced regions labeled - - Note: - Assumes the mid-z slice of the data contains the heart. - - Example: - >>> contrast_labels = segmenter.segment_contrast_agent(preprocessed_image, base_labels) - """ - thoracic_ids = ( - list(self.taxonomy.labels_in_group("heart").keys()) - + list(self.taxonomy.labels_in_group("lung").keys()) - + list(self.taxonomy.labels_in_group("major_vessels").keys()) - + [0] - ) - contrast_ids = list(self.taxonomy.labels_in_group("contrast").keys()) - if len(contrast_ids) == 0: - self.log_warning("No contrast-enhanced regions found in the labelmap") - return labelmap_image - - results_image = self.segment_connected_component( - preprocessed_image, - labelmap_image, - lower_threshold=self.contrast_threshold, - upper_threshold=4000, - labelmap_ids=thoracic_ids, - mask_id=contrast_ids[-1], - use_mid_slice=True, - hole_fill=3, - ) - - return results_image diff --git a/src/physiomotion4d/segment_chest_total_segmentator_with_contrast.py b/src/physiomotion4d/segment_chest_total_segmentator_with_contrast.py new file mode 100644 index 0000000..9d1190a --- /dev/null +++ b/src/physiomotion4d/segment_chest_total_segmentator_with_contrast.py @@ -0,0 +1,225 @@ +"""Module for segmenting contrast-enhanced chest CT images with TotalSegmentator. + +This module provides the SegmentChestTotalSegmentatorWithContrast class, which +extends SegmentChestTotalSegmentator with an additional connected-component +pass that labels contrast-enhanced blood (in the heart, vessels, and lungs) +under a dedicated "contrast" taxonomy group. +""" + +import logging +from typing import Optional + +import itk +import numpy as np + +from .image_tools import ImageTools +from .segment_chest_total_segmentator import SegmentChestTotalSegmentator + + +class SegmentChestTotalSegmentatorWithContrast(SegmentChestTotalSegmentator): + """ + Chest CT segmentation using TotalSegmentator, with contrast-enhanced blood detection. + + Extends :class:`SegmentChestTotalSegmentator` with an additional + connected-component pass that identifies contrast-enhanced blood vessels + and cardiac chambers, labeling them under a ``"contrast"`` taxonomy + group (label id 135). Use this class instead of + :class:`SegmentChestTotalSegmentator` for contrast-enhanced studies. + + Attributes: + contrast_threshold (int): Lower intensity threshold used to detect + contrast-enhanced blood. + + Example: + >>> segmenter = SegmentChestTotalSegmentatorWithContrast() + >>> result = segmenter.segment(ct_image) + >>> labelmap = result['labelmap'] + >>> contrast_mask = result['contrast'] + """ + + def __init__(self, log_level: int | str = logging.INFO): + """Initialize the contrast-enhanced TotalSegmentator-based segmentation. + + Args: + log_level: Logging level (default: logging.INFO) + """ + super().__init__(log_level=log_level) + + self.contrast_threshold: int = 500 + + def _add_extra_taxonomy_groups(self) -> None: + """Register the ``"contrast"`` taxonomy group (label id 135).""" + self.taxonomy.add_organ("contrast", 135, "contrast") + + def postprocess_after_labelmap( + self, input_image: itk.image, labelmap_image: itk.image + ) -> itk.image: + """Run contrast-enhanced blood detection on the labelmap. + + Overrides :meth:`SegmentAnatomyBase.postprocess_after_labelmap`. + + Args: + input_image (itk.image): The original, unpreprocessed input image + labelmap_image (itk.image): The postprocessed segmentation labelmap + + Returns: + itk.image: The labelmap, with contrast-enhanced regions labeled + """ + return self.segment_contrast_agent(input_image, labelmap_image) + + def segment_connected_component( + self, + preprocessed_image: itk.image, + labelmap_image: itk.image, + lower_threshold: int, + upper_threshold: int, + labelmap_ids: Optional[list[int]] = None, + mask_id: int = 0, + use_mid_slice: bool = True, + hole_fill: int = 2, + ) -> itk.image: + """ + Segment connected components based on intensity thresholding. + + Identifies connected regions within intensity thresholds and existing + anatomical masks, then selects the largest component. This is useful + for segmenting structures like contrast-enhanced blood or specific + tissue types. + + Args: + preprocessed_image (itk.image): The preprocessed input image + labelmap_image (itk.image): Existing labelmap to constrain search + lower_threshold (int): Lower intensity threshold + upper_threshold (int): Upper intensity threshold + labelmap_ids (Optional[list[int]]): List of label IDs to search within. + If None, searches within all existing labels + mask_id (int): ID to assign to the segmented component + use_mid_slice (bool): If True, find largest component in middle + slice only; if False, use entire 3D volume + hole_fill (int): Number of pixels to dilate/erode for hole filling + + Returns: + itk.image: Updated labelmap with new component labeled as mask_id + + Example: + >>> # Segment contrast-enhanced blood + >>> updated_labels = segmenter.segment_connected_component( + ... preprocessed_image, labels, 700, 4000, mask_id=135 + ... ) + """ + thresh_image = itk.binary_threshold_image_filter( + Input=preprocessed_image, + LowerThreshold=lower_threshold, + UpperThreshold=upper_threshold, + InsideValue=1, + OutsideValue=0, + ) + thresh_arr = itk.GetArrayFromImage(thresh_image).astype(np.int16) + thresh_image = itk.GetImageFromArray(thresh_arr) + thresh_image.CopyInformation(preprocessed_image) + + label_arr = itk.GetArrayFromImage(labelmap_image) + if labelmap_ids is None: + labelmap_ids = list(self.taxonomy.all_labels().keys()) + label_arr = np.isin(label_arr, labelmap_ids) + label_image = itk.GetImageFromArray(label_arr.astype(np.int16)) + label_image.CopyInformation(labelmap_image) + + connected_component_image = itk.connected_component_image_filter( + Input=thresh_image, + MaskImage=label_image, + ) + + connected_component_arr = itk.GetArrayFromImage(connected_component_image) + if use_mid_slice: + mid_slice = ( + connected_component_image.GetLargestPossibleRegion().GetSize()[2] // 2 + ) + tmp_connected_component_arr = connected_component_arr[mid_slice, :, :] + ids = np.unique(tmp_connected_component_arr) + if len(ids[ids != 0]) > 0: + connected_component_arr = tmp_connected_component_arr + + ids = np.unique(connected_component_arr) + ids = ids[ids != 0] + if ids.size == 0: + self.log_debug( + "segment_connected_component: no connected components found " + "in threshold [%d, %d]; returning labelmap unchanged", + lower_threshold, + upper_threshold, + ) + return labelmap_image + component_sums = [np.sum(connected_component_arr == id) for id in ids] + largest_id = ids[np.argmax(component_sums)] + connected_component_image = itk.binary_threshold_image_filter( + Input=connected_component_image, + LowerThreshold=int(largest_id), + UpperThreshold=int(largest_id), + InsideValue=1, + OutsideValue=0, + ) + image_tools = ImageTools() + connected_component_image = image_tools.binary_dilate_image( + connected_component_image, hole_fill, 1, 0 + ) + connected_component_image = image_tools.binary_erode_image( + connected_component_image, hole_fill, 1, 0 + ) + + labelmap_arr = itk.GetArrayFromImage(labelmap_image) + connected_component_arr = itk.GetArrayFromImage(connected_component_image) + connected_component_mask = connected_component_arr > 0 + mask = label_arr & connected_component_mask + labelmap_arr = np.where(mask, mask_id, labelmap_arr) + results_image = itk.GetImageFromArray(labelmap_arr.astype(np.uint8)) + results_image.CopyInformation(preprocessed_image) + + return results_image + + def segment_contrast_agent( + self, preprocessed_image: itk.image, labelmap_image: itk.image + ) -> itk.image: + """ + Include contrast-enhanced blood in the labelmap. + + Segments high-intensity regions corresponding to contrast-enhanced + blood vessels and cardiac chambers. Uses connected component analysis + focused on the middle slice where the heart is typically located. + + Args: + preprocessed_image (itk.image): The preprocessed CT image + labelmap_image (itk.image): Existing segmentation labelmap + + Returns: + itk.image: Updated labelmap with contrast-enhanced regions labeled + + Note: + Assumes the mid-z slice of the data contains the heart. + + Example: + >>> contrast_labels = segmenter.segment_contrast_agent(preprocessed_image, base_labels) + """ + thoracic_ids = ( + list(self.taxonomy.labels_in_group("heart").keys()) + + list(self.taxonomy.labels_in_group("lung").keys()) + + list(self.taxonomy.labels_in_group("major_vessels").keys()) + + [0] + ) + contrast_ids = list(self.taxonomy.labels_in_group("contrast").keys()) + if len(contrast_ids) == 0: + self.log_warning("No contrast-enhanced regions found in the labelmap") + return labelmap_image + + results_image = self.segment_connected_component( + preprocessed_image, + labelmap_image, + lower_threshold=self.contrast_threshold, + upper_threshold=4000, + labelmap_ids=thoracic_ids, + mask_id=contrast_ids[-1], + use_mid_slice=True, + hole_fill=3, + ) + + return results_image diff --git a/src/physiomotion4d/workflow_convert_image_to_usd.py b/src/physiomotion4d/workflow_convert_image_to_usd.py index 3fc8756..b537d96 100644 --- a/src/physiomotion4d/workflow_convert_image_to_usd.py +++ b/src/physiomotion4d/workflow_convert_image_to_usd.py @@ -22,7 +22,9 @@ from .register_images_base import RegisterImagesBase from .register_images_icon import RegisterImagesICON from .segment_anatomy_base import SegmentAnatomyBase -from .segment_chest_total_segmentator import SegmentChestTotalSegmentator +from .segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from .transform_tools import TransformTools from .usd_anatomy_tools import USDAnatomyTools @@ -38,7 +40,7 @@ class WorkflowConvertImageToUSD(PhysioMotion4DBase): pre-configured :class:`SegmentAnatomyBase` / :class:`RegisterImagesBase` instance. Configure backend-specific parameters (iteration counts, trim_branches, mass preservation, etc.) on the instance before passing - it in. Defaults to :class:`SegmentChestTotalSegmentator` / + it in. Defaults to :class:`SegmentChestTotalSegmentatorWithContrast` / :class:`RegisterImagesICON` when omitted. """ @@ -66,7 +68,7 @@ def __init__( output_directory (str): Directory path where output files will be stored segmentation_method (Optional[SegmentAnatomyBase]): Segmentation backend instance. Defaults to a new - :class:`SegmentChestTotalSegmentator` when None. + :class:`SegmentChestTotalSegmentatorWithContrast` when None. registration_method (Optional[RegisterImagesBase]): Registration backend instance. Defaults to a new :class:`RegisterImagesICON` when None. A caller-supplied instance is mutated (fixed @@ -103,9 +105,9 @@ def __init__( dict[str, dict[str, Union[itk.Transform, float]]] ] = [] if segmentation_method is None: - segmentation_method = SegmentChestTotalSegmentator(log_level=log_level) - segmentation_method.contrast_threshold = 500 - segmentation_method.set_contrast_enhanced_study(True) + segmentation_method = SegmentChestTotalSegmentatorWithContrast( + log_level=log_level + ) elif not isinstance(segmentation_method, SegmentAnatomyBase): raise TypeError( "segmentation_method must be a SegmentAnatomyBase instance or None" diff --git a/src/physiomotion4d/workflow_convert_image_to_vtk.py b/src/physiomotion4d/workflow_convert_image_to_vtk.py index 5cff620..4e240b5 100644 --- a/src/physiomotion4d/workflow_convert_image_to_vtk.py +++ b/src/physiomotion4d/workflow_convert_image_to_vtk.py @@ -9,12 +9,15 @@ Typical usage:: import itk - from physiomotion4d import SegmentChestTotalSegmentator, WorkflowConvertImageToVTK + from physiomotion4d import ( + SegmentChestTotalSegmentatorWithContrast, + WorkflowConvertImageToVTK, + ) ct = itk.imread('chest_ct.nii.gz') - segmenter = SegmentChestTotalSegmentator() + segmenter = SegmentChestTotalSegmentatorWithContrast() workflow = WorkflowConvertImageToVTK(segmentation_method=segmenter) - result = workflow.run_workflow(ct, contrast_enhanced_study=True) + result = workflow.run_workflow(ct) # Combined single-file output (default) WorkflowConvertImageToVTK.save_combined_surface(result['surfaces'], './out', prefix='patient') @@ -36,7 +39,9 @@ from .contour_tools import ContourTools from .physiomotion4d_base import PhysioMotion4DBase from .segment_anatomy_base import SegmentAnatomyBase -from .segment_chest_total_segmentator import SegmentChestTotalSegmentator +from .segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from .usd_anatomy_tools import USDAnatomyTools @@ -45,10 +50,11 @@ class WorkflowConvertImageToVTK(PhysioMotion4DBase): ``segmentation_method`` accepts a pre-configured :class:`SegmentAnatomyBase` instance (e.g. :class:`SegmentChestTotalSegmentator`, - :class:`SegmentHeartSimpleware`, or :class:`SegmentHeartSimplewareTrimmedBranches` - for cardiac-only segmentation with pulmonary/great-vessel branches - trimmed). Defaults to a new :class:`SegmentChestTotalSegmentator` when - omitted. + :class:`SegmentChestTotalSegmentatorWithContrast` for contrast-enhanced + studies, :class:`SegmentHeartSimpleware`, or + :class:`SegmentHeartSimplewareTrimmedBranches` for cardiac-only + segmentation with pulmonary/great-vessel branches trimmed). Defaults to + a new :class:`SegmentChestTotalSegmentator` when omitted. **Output anatomy groups** @@ -85,7 +91,7 @@ def __init__( Args: segmentation_method: Segmentation backend instance. Defaults to - a new :class:`SegmentChestTotalSegmentator` when None. + a new :class:`SegmentChestTotalSegmentatorWithContrast` when None. log_level: Logging level. Default: ``logging.INFO``. Raises: @@ -95,7 +101,9 @@ def __init__( super().__init__(class_name=self.__class__.__name__, log_level=log_level) if segmentation_method is None: - segmentation_method = SegmentChestTotalSegmentator(log_level=log_level) + segmentation_method = SegmentChestTotalSegmentatorWithContrast( + log_level=log_level + ) elif not isinstance(segmentation_method, SegmentAnatomyBase): raise TypeError( "segmentation_method must be a SegmentAnatomyBase instance or None" @@ -206,19 +214,12 @@ def _extract_mesh(self, mask_image: Any) -> Optional[pv.UnstructuredGrid]: def run_workflow( self, input_image: Any, - contrast_enhanced_study: bool = False, anatomy_groups: Optional[list[str]] = None, ) -> dict[str, Any]: """Segment the CT image and extract per-anatomy-group VTK objects. Args: input_image: Input 3D CT image (``itk.Image``). - contrast_enhanced_study: If ``True``, an additional connected-component - pass identifies contrast-enhanced blood. Default: ``False``. - Only supported by segmenters that implement - ``set_contrast_enhanced_study`` (e.g. - :class:`SegmentChestTotalSegmentator`); requesting ``True`` - with any other segmenter raises :class:`ValueError`. anatomy_groups: Subset of anatomy groups to process. ``None`` (default) processes all non-empty groups. Valid names are given by :attr:`ANATOMY_GROUPS`, derived from the active segmenter's @@ -255,18 +256,6 @@ def run_workflow( self.log_info("Running segmenter: %s", type(self._segmenter).__name__) self.log_section("Running segmentation") - set_contrast_enhanced_study = getattr( - self._segmenter, "set_contrast_enhanced_study", None - ) - if contrast_enhanced_study and set_contrast_enhanced_study is None: - raise ValueError( - "contrast_enhanced_study=True requires a segmenter that " - "implements set_contrast_enhanced_study (e.g. " - "SegmentChestTotalSegmentator); " - f"got {type(self._segmenter).__name__}" - ) - if set_contrast_enhanced_study is not None: - set_contrast_enhanced_study(contrast_enhanced_study) seg_result: dict[str, Any] = self._segmenter.segment(input_image) # Extract VTK objects per anatomy group diff --git a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py index ca8c310..07de1b8 100644 --- a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py @@ -192,7 +192,6 @@ def __init__( ) patient_models_data = convert_image_to_vtk.run_workflow( input_image=patient_image, - contrast_enhanced_study=False, anatomy_groups=["heart"], ) patient_models = [patient_models_data["meshes"]["heart"]] diff --git a/tests/conftest.py b/tests/conftest.py index b7c9ce0..f6a2457 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,9 @@ from physiomotion4d.register_images_greedy import RegisterImagesGreedy from physiomotion4d.register_images_icon import RegisterImagesICON from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator +from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from physiomotion4d.segment_heart_simpleware import SegmentHeartSimpleware from physiomotion4d.transform_tools import TransformTools @@ -601,6 +604,14 @@ def segmenter_total_segmentator() -> SegmentChestTotalSegmentator: return SegmentChestTotalSegmentator() +@pytest.fixture(scope="session") +def segmenter_total_segmentator_with_contrast() -> ( + SegmentChestTotalSegmentatorWithContrast +): + """Create a SegmentChestTotalSegmentatorWithContrast instance.""" + return SegmentChestTotalSegmentatorWithContrast() + + @pytest.fixture(scope="session") def segmenter_simpleware() -> SegmentHeartSimpleware: """Create a SegmentHeartSimpleware instance.""" diff --git a/tests/test_segment_chest_total_segmentator.py b/tests/test_segment_chest_total_segmentator.py index a9ef7c0..ca4d8df 100644 --- a/tests/test_segment_chest_total_segmentator.py +++ b/tests/test_segment_chest_total_segmentator.py @@ -14,6 +14,9 @@ import pytest from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator +from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) @pytest.mark.requires_gpu @@ -80,7 +83,6 @@ def test_segment_single_image( "bone", "soft_tissue", "other", - "contrast", ] for key in expected_keys: assert key in result, f"Missing key '{key}' in result" @@ -182,32 +184,29 @@ def test_anatomy_group_masks( def test_contrast_detection( self, segmenter_total_segmentator: SegmentChestTotalSegmentator, + segmenter_total_segmentator_with_contrast: ( + SegmentChestTotalSegmentatorWithContrast + ), test_images: list[Any], ) -> None: """Test contrast detection functionality.""" input_image = test_images[0] - # Test without contrast - segmenter_total_segmentator.set_contrast_enhanced_study(False) + # Plain segmenter has no "contrast" group result_no_contrast = segmenter_total_segmentator.segment(input_image) - contrast_mask_no = result_no_contrast["contrast"] + assert "contrast" not in result_no_contrast - # Test with contrast flag - segmenter_total_segmentator.set_contrast_enhanced_study(True) - result_with_contrast = segmenter_total_segmentator.segment(input_image) + # Contrast-enhanced segmenter adds a "contrast" group + result_with_contrast = segmenter_total_segmentator_with_contrast.segment( + input_image + ) contrast_mask_yes = result_with_contrast["contrast"] - segmenter_total_segmentator.set_contrast_enhanced_study(False) - - # Both should return valid masks - assert contrast_mask_no is not None, "Contrast mask (no flag) is None" assert contrast_mask_yes is not None, "Contrast mask (with flag) is None" print("\nContrast detection tested") - contrast_arr_no = itk.array_from_image(contrast_mask_no) contrast_arr_yes = itk.array_from_image(contrast_mask_yes) - print(f" Without contrast flag: {np.sum(contrast_arr_no > 0)} voxels") - print(f" With contrast flag: {np.sum(contrast_arr_yes > 0)} voxels") + print(f" Contrast voxels: {np.sum(contrast_arr_yes > 0)}") def test_preprocessing( self, diff --git a/tests/test_workflow_convert_image_to_usd.py b/tests/test_workflow_convert_image_to_usd.py index d7d9ded..080e166 100644 --- a/tests/test_workflow_convert_image_to_usd.py +++ b/tests/test_workflow_convert_image_to_usd.py @@ -13,7 +13,9 @@ from physiomotion4d.register_images_base import RegisterImagesBase from physiomotion4d.register_images_icon import RegisterImagesICON -from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator +from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, +) from physiomotion4d.workflow_convert_image_to_usd import WorkflowConvertImageToUSD @@ -24,7 +26,7 @@ def _small_image() -> itk.Image: def test_default_segmentation_and_registration_methods(tmp_path: Path) -> None: """Omitting segmentation_method/registration_method defaults to - SegmentChestTotalSegmentator (contrast_threshold=500) and + SegmentChestTotalSegmentatorWithContrast (contrast_threshold=500) and RegisterImagesICON, matching this workflow's documented defaults.""" reference_image = _small_image() workflow = WorkflowConvertImageToUSD( @@ -35,7 +37,7 @@ def test_default_segmentation_and_registration_methods(tmp_path: Path) -> None: log_level=logging.CRITICAL, ) - assert isinstance(workflow.segmenter, SegmentChestTotalSegmentator) + assert isinstance(workflow.segmenter, SegmentChestTotalSegmentatorWithContrast) assert workflow.segmenter.contrast_threshold == 500 assert isinstance(workflow.registrar, RegisterImagesICON) @@ -73,7 +75,8 @@ def test_caller_supplied_instances_are_used_as_is(tmp_path: Path) -> None: (beyond the documented shared setters): the workflow must not apply its default-only contrast_threshold=500 tuning to a caller-supplied segmenter.""" reference_image = _small_image() - segmenter = SegmentChestTotalSegmentator() + segmenter = SegmentChestTotalSegmentatorWithContrast() + segmenter.contrast_threshold = 800 original_contrast_threshold = segmenter.contrast_threshold registrar: RegisterImagesBase = RegisterImagesICON() @@ -112,9 +115,8 @@ def test_workflow_convert_image_to_usd_default_operation( log_level=logging.CRITICAL, ) - assert isinstance(workflow.segmenter, SegmentChestTotalSegmentator) + assert isinstance(workflow.segmenter, SegmentChestTotalSegmentatorWithContrast) assert workflow.segmenter.contrast_threshold == 500 - assert workflow.segmenter.contrast_enhanced_study assert isinstance(workflow.registrar, RegisterImagesICON) workflow.registrar.set_number_of_iterations(2) diff --git a/tests/test_workflow_fit_statistical_model_to_patient.py b/tests/test_workflow_fit_statistical_model_to_patient.py index 5c1c0a9..5cf28ce 100644 --- a/tests/test_workflow_fit_statistical_model_to_patient.py +++ b/tests/test_workflow_fit_statistical_model_to_patient.py @@ -120,7 +120,6 @@ def run_workflow(self, **kwargs: Any) -> dict[str, Any]: SegmentHeartSimplewareTrimmedBranches, ) assert captured["run_kwargs"]["anatomy_groups"] == ["heart"] - assert captured["run_kwargs"]["contrast_enhanced_study"] is False assert workflow.patient_models == [heart_mesh] diff --git a/tutorials/tutorial_02_ct_to_vtk.py b/tutorials/tutorial_02_ct_to_vtk.py index fe946a4..895edde 100644 --- a/tutorials/tutorial_02_ct_to_vtk.py +++ b/tutorials/tutorial_02_ct_to_vtk.py @@ -23,8 +23,8 @@ import itk import pyvista as pv -from physiomotion4d.segment_chest_total_segmentator import ( - SegmentChestTotalSegmentator, +from physiomotion4d.segment_chest_total_segmentator_with_contrast import ( + SegmentChestTotalSegmentatorWithContrast, ) from physiomotion4d.test_tools import TestTools from physiomotion4d.workflow_convert_image_to_vtk import WorkflowConvertImageToVTK @@ -71,16 +71,15 @@ # %% # Workflow initialization workflow = WorkflowConvertImageToVTK( - segmentation_method=SegmentChestTotalSegmentator(log_level=log_level), + segmentation_method=SegmentChestTotalSegmentatorWithContrast( + log_level=log_level + ), log_level=log_level, ) # %% # Workflow execution - result = workflow.run_workflow( - input_image=ct_image, - contrast_enhanced_study=True, - ) + result = workflow.run_workflow(input_image=ct_image) # %% # Result saving