diff --git a/.claude/.gitignore b/.claude/.gitignore new file mode 100644 index 0000000..93c0f73 --- /dev/null +++ b/.claude/.gitignore @@ -0,0 +1 @@ +settings.local.json diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 61baeef..cce1e5d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -107,6 +107,7 @@ jobs: uses: codecov/codecov-action@v4 if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.10' with: + token: ${{ secrets.CODECOV_TOKEN }} file: ./coverage.xml flags: unittests name: codecov-unit-${{ matrix.os }}-py${{ matrix.python-version }} @@ -233,6 +234,7 @@ jobs: - name: Upload coverage to Codecov uses: codecov/codecov-action@v4 with: + token: ${{ secrets.CODECOV_TOKEN }} file: ./coverage.xml flags: integration-tests name: codecov-integration @@ -317,6 +319,7 @@ jobs: - name: Upload coverage to Codecov uses: codecov/codecov-action@v4 with: + token: ${{ secrets.CODECOV_TOKEN }} file: ./coverage.xml flags: gpu-tests name: codecov-gpu @@ -380,7 +383,6 @@ jobs: # - tests/test_register_images_icon.py (requires CUDA for ICON) # - tests/test_transform_tools.py (depends on slow registration tests) # - tests/test_segment_chest_total_segmentator.py (requires CUDA for TotalSegmentator) -# - tests/test_segment_chest_vista_3d.py (requires CUDA for VISTA-3D, 20GB+ RAM) # # Experiment tests (EXTREMELY SLOW - hours to complete): # - tests/test_experiments.py (runs all notebooks in experiments/ subdirectories) diff --git a/.github/workflows/test-slow.yml b/.github/workflows/test-slow.yml index 3ce4e49..80ebcc7 100644 --- a/.github/workflows/test-slow.yml +++ b/.github/workflows/test-slow.yml @@ -80,6 +80,7 @@ jobs: - name: Upload coverage to Codecov uses: codecov/codecov-action@v4 with: + token: ${{ secrets.CODECOV_TOKEN }} file: ./coverage.xml flags: slow-tests-gpu name: codecov-slow-gpu diff --git a/CLAUDE.md b/CLAUDE.md index e577399..e587dd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -28,7 +28,6 @@ py -m pytest tests/test_contour_tools.py::test_extract_surface -v # Skip GPU-dependent tests py -m pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py \ - --ignore=tests/test_segment_chest_vista_3d.py \ --ignore=tests/test_register_images_icon.py # With coverage diff --git a/README.md b/README.md index 83cab7d..84c9887 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ PhysioMotion4D is a comprehensive medical imaging package that converts 4D CT sc ## 🚀 Key Features - **Complete 4D Medical Imaging Pipeline**: End-to-end processing from 4D CT data to animated USD models -- **Multiple AI Segmentation Methods**: TotalSegmentator, VISTA-3D, and ensemble approaches +- **Multiple AI Segmentation Methods**: TotalSegmentator and ensemble approaches - **Deep Learning Registration**: GPU-accelerated image registration using Icon algorithm - **NVIDIA Omniverse Integration**: Direct USD file export for medical visualization - **Physiological Motion Analysis**: Capture and visualize cardiac and respiratory motion @@ -106,8 +106,6 @@ print(f"PhysioMotion4D version: {physiomotion4d.__version__}") - `WorkflowFitStatisticalModelToPatient`: Model-to-patient registration workflow - **Segmentation Classes**: Multiple AI-based chest segmentation implementations - `SegmentChestTotalSegmentator`: TotalSegmentator-based segmentation - - `SegmentChestVista3D`: VISTA-3D model-based segmentation - - `SegmentChestVista3DNIM`: NVIDIA NIM version of VISTA-3D - `SegmentChestEnsemble`: Ensemble segmentation combining multiple methods - `SegmentAnatomyBase`: Base class for custom segmentation methods - **Registration Classes**: Multiple registration methods for different use cases @@ -142,7 +140,7 @@ print(f"PhysioMotion4D version: {physiomotion4d.__version__}") - **AI/ML**: PyTorch, CuPy (CUDA 13 default; CUDA 12 via `[cuda12]` extra), transformers, MONAI - **Registration**: icon-registration, unigradicon - **Visualization**: USD-core, PyVista -- **Segmentation**: TotalSegmentator, VISTA-3D models +- **Segmentation**: TotalSegmentator ## 🎯 Quick Start @@ -262,11 +260,11 @@ registered_mesh = workflow.run_workflow() ### Custom Segmentation ```python -from physiomotion4d import SegmentChestVista3D +from physiomotion4d import SegmentChestTotalSegmentator import itk -# Initialize VISTA-3D segmentation -segmenter = SegmentChestVista3D() +# Initialize TotalSegmentator segmentation +segmenter = SegmentChestTotalSegmentator() # Load and segment image image = itk.imread("chest_ct.nrrd") @@ -578,14 +576,12 @@ pytest tests/test_register_images_greedy.py -v # Greedy registration pytest tests/test_register_images_icon.py -v # Icon registration pytest tests/test_register_time_series_images.py -v # Time series registration pytest tests/test_segment_chest_total_segmentator.py -v # TotalSegmentator -pytest tests/test_segment_chest_vista_3d.py -v # VISTA-3D segmentation pytest tests/test_contour_tools.py -v # Mesh and contour tools pytest tests/test_image_tools.py -v # Image processing utilities pytest tests/test_transform_tools.py -v # Transform operations # Skip GPU-dependent tests (segmentation and registration) pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py \ - --ignore=tests/test_segment_chest_vista_3d.py \ --ignore=tests/test_register_images_icon.py # Run with coverage report @@ -594,7 +590,7 @@ pytest tests/ --cov=src/physiomotion4d --cov-report=html **Test Categories:** - **Data Pipeline**: Download, conversion, and preprocessing -- **Segmentation**: TotalSegmentator and VISTA-3D (GPU required) +- **Segmentation**: TotalSegmentator (GPU required) - **Registration**: ANTs, Icon, and time series methods (slow, ~5-10 min) - **Geometry & Visualization**: Contour tools, transform tools, VTK to USD - **USD Utilities**: Merging, time preservation, material handling @@ -767,7 +763,7 @@ This project is licensed under the Apache 2.0 License - see the LICENSE file for - **NVIDIA Omniverse** team for USD format and visualization platform - **MONAI** community for medical imaging AI tools - **DirLab** for providing the 4D-CT benchmark datasets -- **TotalSegmentator** and **VISTA-3D** teams for segmentation models +- **TotalSegmentator** team for segmentation models - **Icon Registration** team for deep learning registration methods ## 📞 Support diff --git a/docs/API_MAP.md b/docs/API_MAP.md index 07d6aba..259d799 100644 --- a/docs/API_MAP.md +++ b/docs/API_MAP.md @@ -31,10 +31,6 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def transform_contours(contours, transform_filenames, frame_indices, base_name, output_dir)` (line 31) - `def convert_contours(base_name, output_dir, project_name, compute_normals=False)` (line 49) -## experiments/Heart-GatedCT_To_USD/test_vista3d_inMem.py - -- `def vista3d_inference_from_itk(itk_image, label_prompt=None, points=None, point_labels=None, device=None, bundle_path=None, model_cache_dir=None)` (line 8) - ## experiments/Lung-GatedCT_To_USD/0-register_dirlab_4dct.py - `def dilate_mask(mask, dilation)` (line 30) @@ -137,56 +133,6 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def convert_array_to_image_of_vectors(self, arr_data, reference_image, ptype=itk.D)` (line 218): Convert a numpy array to an ITK image of vector type. - `def flip_image(self, in_image, in_mask=None, flip_x=False, flip_y=False, flip_z=False, flip_and_make_identity=False)` (line 249): Flip the image and mask. -## src/physiomotion4d/network_weights/vista3d/hugging_face_pipeline.py - -- **class HuggingFacePipelineHelper** (line 7) - - `def __init__(self, pipeline_name='vista3d')` (line 9) - - `def get_pipeline(self)` (line 18) - - `def init_pipeline(self, pretrained_model_name_or_path, **kwargs)` (line 30) - -## src/physiomotion4d/network_weights/vista3d/scripts/early_stop_score_function.py - -- `def score_function(engine)` (line 7) - -## src/physiomotion4d/network_weights/vista3d/scripts/evaluator.py - -- **class Vista3dEvaluator** (line 39): Supervised detection evaluation method with image and label, inherits from ``SupervisedEvaluator`` and ``Workflow``. - - `def __init__(self, device, val_data_loader, network, epoch_length=None, non_blocking=False, prepare_batch=default_prepare_batch, iteration_update=None, inferer=None, postprocessing=None, key_val_metric=None, additional_metrics=None, metric_cmp_fn=default_metric_cmp_fn, val_handlers=None, amp=False, mode=ForwardMode.EVAL, event_names=None, event_to_attr=None, decollate=True, to_kwargs=None, amp_kwargs=None, hyper_kwargs=None)` (line 85) - - `def transform_points(self, point, affine)` (line 137): transform point to the coordinates of the transformed image - - `def check_prompts_format(self, label_prompt, points, point_labels)` (line 148): check the format of user prompts - -## src/physiomotion4d/network_weights/vista3d/scripts/inferer.py - -- **class Vista3dInferer** (line 21): Vista3D Inferer - - `def __init__(self, roi_size, overlap, use_point_window=False, sw_batch_size=1)` (line 30) - -## src/physiomotion4d/network_weights/vista3d/scripts/trainer.py - -- **class Vista3dTrainer** (line 39): Supervised detection training method with image and label, inherits from ``Trainer`` and ``Workflow``. - - `def __init__(self, device, max_epochs, train_data_loader, network, optimizer, loss_function, epoch_length=None, non_blocking=False, prepare_batch=default_prepare_batch, iteration_update=None, inferer=None, postprocessing=None, key_train_metric=None, additional_metrics=None, metric_cmp_fn=default_metric_cmp_fn, train_handlers=None, amp=False, event_names=None, event_to_attr=None, decollate=True, optim_set_to_none=False, to_kwargs=None, amp_kwargs=None, hyper_kwargs=None)` (line 88) - -## src/physiomotion4d/network_weights/vista3d/vista3d_config.py - -- **class VISTA3DConfig** (line 4): Configuration class for vista3d - - `def __init__(self, encoder_embed_dim=48, input_channels=1, **kwargs)` (line 9): Set the hyperparameters for the VISTA3D model. - -## src/physiomotion4d/network_weights/vista3d/vista3d_model.py - -- **class VISTA3DModel** (line 9): VISTA3D model for hugging face - - `def __init__(self, config)` (line 14) - - `def forward(self, input)` (line 22) -- `def register_my_model()` (line 26): Utility function to register VISTA3D model so that it can be instantiate by the AutoModel function. - -## src/physiomotion4d/network_weights/vista3d/vista3d_pipeline.py - -- **class VISTA3DPipeline** (line 48): Define the VISTA3D pipeline. - - `def __init__(self, model, **kwargs)` (line 80) - - `def check_prompts_format(self, label_prompt, points, point_labels)` (line 213): check the format of user prompts - - `def transform_points(self, point, affine)` (line 287): transform point to the coordinates of the transformed image - - `def preprocess(self, inputs, **kwargs)` (line 298) - - `def postprocess(self, outputs, **kwargs)` (line 435) -- `def register_simple_pipeline()` (line 454) - ## src/physiomotion4d/notebook_utils.py - `def running_as_test()` (line 11): True when the notebook is run as a test (e.g. by pytest experiment tests). @@ -321,10 +267,8 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ ## src/physiomotion4d/segment_chest_ensemble.py -- **class SegmentChestEnsemble** (line 20): A class that inherits from physioSegmentChest and implements the - - `def __init__(self, log_level=logging.INFO)` (line 26): Initialize the vista3d class. - - `def ensemble_segmentation(self, labelmap_vista, labelmap_totseg)` (line 309): Combine two segmentation results using label mapping and priority rules. - - `def segmentation_method(self, preprocessed_image)` (line 398): Run VISTA3D on the preprocessed image and return result. +- **class SegmentChestEnsemble** (line 12): Ensemble chest CT segmentation. + - `def __init__(self, log_level=logging.INFO)` (line 29) ## src/physiomotion4d/segment_chest_total_segmentator.py @@ -332,22 +276,6 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def __init__(self, log_level=logging.INFO)` (line 57): Initialize the TotalSegmentator-based chest segmentation. - `def segmentation_method(self, preprocessed_image)` (line 199): Run TotalSegmentator on the preprocessed image and return result. -## src/physiomotion4d/segment_chest_vista_3d.py - -- **class SegmentChestVista3D** (line 32): Chest CT segmentation using NVIDIA VISTA-3D foundational model. - - `def __init__(self, log_level=logging.INFO)` (line 68): Initialize the VISTA-3D based chest segmentation. - - `def set_label_prompt(self, label_prompt)` (line 232): Set specific anatomical structure labels to segment. - - `def set_whole_image_segmentation(self)` (line 253): Configure for automatic whole-image segmentation. - - `def segment_soft_tissue(self, preprocessed_image, labelmap_image)` (line 269): Add soft tissue segmentation to fill gaps in VISTA-3D output. - - `def preprocess_input(self, input_image)` (line 302): Preprocess the input image for VISTA-3D segmentation. - - `def segmentation_method(self, preprocessed_image)` (line 329): Run VISTA-3D segmentation on the preprocessed image. - -## src/physiomotion4d/segment_chest_vista_3d_nim.py - -- **class SegmentChestVista3DNIM** (line 25): A class that inherits from physioSegmentChest and implements the - - `def __init__(self, log_level=logging.INFO)` (line 31): Initialize the vista3d class. - - `def segmentation_method(self, preprocessed_image)` (line 45): Run VISTA3D on the preprocessed image using the NIM and return result. - ## src/physiomotion4d/segment_heart_simpleware.py - **class SegmentHeartSimpleware** (line 23): Heart CT segmentation using Simpleware Medical's ASCardio module. @@ -490,13 +418,13 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ ## src/physiomotion4d/workflow_convert_ct_to_vtk.py -- **class WorkflowConvertCTToVTK** (line 59): Segment a CT image and produce per-anatomy-group VTK surfaces and meshes. - - `def __init__(self, segmentation_method='total_segmentator', log_level=logging.INFO)` (line 100): Initialize the workflow. - - `def run_workflow(self, input_image, contrast_enhanced_study=False, anatomy_groups=None)` (line 247): Segment the CT image and extract per-anatomy-group VTK objects. - - `def save_surfaces(surfaces, output_dir, prefix='')` (line 350): Save each group surface to its own VTP file. - - `def save_meshes(meshes, output_dir, prefix='')` (line 377): Save each group voxel mesh to its own VTU file. - - `def save_combined_surface(surfaces, output_dir, prefix='')` (line 403): Merge all group surfaces into a single VTP file. - - `def save_combined_mesh(meshes, output_dir, prefix='')` (line 438): Merge all group meshes into a single VTU file. +- **class WorkflowConvertCTToVTK** (line 58): Segment a CT image and produce per-anatomy-group VTK surfaces and meshes. + - `def __init__(self, segmentation_method='total_segmentator', log_level=logging.INFO)` (line 98): Initialize the workflow. + - `def run_workflow(self, input_image, contrast_enhanced_study=False, anatomy_groups=None)` (line 240): Segment the CT image and extract per-anatomy-group VTK objects. + - `def save_surfaces(surfaces, output_dir, prefix='')` (line 343): Save each group surface to its own VTP file. + - `def save_meshes(meshes, output_dir, prefix='')` (line 370): Save each group voxel mesh to its own VTU file. + - `def save_combined_surface(surfaces, output_dir, prefix='')` (line 396): Merge all group surfaces into a single VTP file. + - `def save_combined_mesh(meshes, output_dir, prefix='')` (line 431): Merge all group meshes into a single VTU file. ## src/physiomotion4d/workflow_convert_heart_gated_ct_to_usd.py @@ -558,20 +486,19 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def pytest_terminal_summary(terminalreporter, exitstatus, config)` (line 119): Print comprehensive test timing report after all tests complete. - `def test_directories()` (line 253): Set up test directories for data and results. - `def download_truncal_valve_data(test_directories)` (line 268): Download TruncalValve 4D CT data. -- `def converted_3d_images(download_truncal_valve_data, test_directories)` (line 310): Convert 4D NRRD to 3D time series and return slice files. -- `def test_images(converted_3d_images)` (line 338): Load time points from the converted 3D data for testing. -- `def segmenter_total_segmentator()` (line 379): Create a SegmentChestTotalSegmentator instance. -- `def segmenter_vista_3d()` (line 385): Create a SegmentChestVista3D instance. -- `def segmenter_simpleware()` (line 391): Create a SegmentHeartSimpleware instance. -- `def heart_simpleware_image_path()` (line 397): Path to cardiac CT image used by experiments/Heart-Simpleware_Segmentation notebook. -- `def heart_simpleware_image(heart_simpleware_image_path)` (line 416): Load cardiac CT image for SegmentHeartSimpleware tests (same as notebook). -- `def segmentation_results(segmenter_total_segmentator, test_images, test_directories)` (line 422): Get or create segmentation results using TotalSegmentator. -- `def contour_tools()` (line 483): Create a ContourTools instance. -- `def registrar_ants()` (line 494): Create a RegisterImagesANTs instance. -- `def registrar_greedy()` (line 500): Create a RegisterImagesGreedy instance. -- `def registrar_icon()` (line 506): Create a RegisterImagesICON instance. -- `def ants_registration_results(registrar_ants, test_images, test_directories)` (line 512): Perform ANTs registration and return results. -- `def transform_tools()` (line 565): Create a TransformTools instance. +- `def converted_3d_images(download_truncal_valve_data, test_directories)` (line 315): Convert 4D NRRD to 3D time series and return slice files. +- `def test_images(converted_3d_images)` (line 343): Load time points from the converted 3D data for testing. +- `def segmenter_total_segmentator()` (line 384): Create a SegmentChestTotalSegmentator instance. +- `def segmenter_simpleware()` (line 390): Create a SegmentHeartSimpleware instance. +- `def heart_simpleware_image_path()` (line 396): Path to cardiac CT image used by experiments/Heart-Simpleware_Segmentation notebook. +- `def heart_simpleware_image(heart_simpleware_image_path)` (line 415): Load cardiac CT image for SegmentHeartSimpleware tests (same as notebook). +- `def segmentation_results(segmenter_total_segmentator, test_images, test_directories)` (line 421): Get or create segmentation results using TotalSegmentator. +- `def contour_tools()` (line 482): Create a ContourTools instance. +- `def registrar_ants()` (line 493): Create a RegisterImagesANTs instance. +- `def registrar_greedy()` (line 499): Create a RegisterImagesGreedy instance. +- `def registrar_icon()` (line 505): Create a RegisterImagesICON instance. +- `def ants_registration_results(registrar_ants, test_images, test_directories)` (line 511): Perform ANTs registration and return results. +- `def transform_tools()` (line 564): Create a TransformTools instance. ## tests/test_contour_tools.py @@ -625,12 +552,12 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def test_experiment_reconstruct_4dct()` (line 331): Test Reconstruct4DCT experiment scripts. - `def test_experiment_heart_vtk_series_to_usd()` (line 350): Test Heart-VTKSeries_To_USD experiment scripts. - `def test_experiment_heart_gated_ct_to_usd()` (line 371): Test Heart-GatedCT_To_USD experiment scripts. -- `def test_experiment_convert_vtk_to_usd()` (line 397): Test Convert_VTK_To_USD experiment scripts. -- `def test_experiment_create_statistical_model()` (line 417): Test Heart-Create_Statistical_Model experiment scripts. -- `def test_experiment_heart_statistical_model_to_patient()` (line 442): Test Heart-Statistical_Model_To_Patient experiment scripts. -- `def test_experiment_lung_gated_ct_to_usd()` (line 477): Test Lung-GatedCT_To_USD experiment scripts. -- `def test_experiment_structure()` (line 522): Validate the structure of the experiments directory. -- `def test_list_scripts_in_subdir(subdir_name)` (line 576): List all scripts in each experiment subdirectory. +- `def test_experiment_convert_vtk_to_usd()` (line 395): Test Convert_VTK_To_USD experiment scripts. +- `def test_experiment_create_statistical_model()` (line 415): Test Heart-Create_Statistical_Model experiment scripts. +- `def test_experiment_heart_statistical_model_to_patient()` (line 440): Test Heart-Statistical_Model_To_Patient experiment scripts. +- `def test_experiment_lung_gated_ct_to_usd()` (line 475): Test Lung-GatedCT_To_USD experiment scripts. +- `def test_experiment_structure()` (line 520): Validate the structure of the experiments directory. +- `def test_list_scripts_in_subdir(subdir_name)` (line 574): List all scripts in each experiment subdirectory. ## tests/test_image_tools.py @@ -735,19 +662,6 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def test_preprocessing(self, segmenter_total_segmentator, test_images)` (line 204): Test preprocessing functionality. - `def test_postprocessing(self, segmenter_total_segmentator, test_images)` (line 224): Test postprocessing functionality. -## tests/test_segment_chest_vista_3d.py - -- **class TestSegmentChestVista3D** (line 16): Test suite for VISTA-3D chest CT segmentation. - - `def test_segmenter_initialization(self, segmenter_vista_3d)` (line 19): Test that SegmentChestVista3D initializes correctly. - - `def test_segment_single_image(self, segmenter_vista_3d, test_images, test_directories)` (line 51): Test automatic segmentation on a single time point. - - `def test_segment_multiple_images(self, segmenter_vista_3d, test_images, test_directories)` (line 107): Test automatic segmentation on two time points. - - `def test_anatomy_group_masks(self, segmenter_vista_3d, test_images)` (line 138): Test that anatomy group masks are created correctly. - - `def test_label_prompt_segmentation(self, segmenter_vista_3d, test_images, test_directories)` (line 177): Test segmentation with specific label prompts. - - `def test_contrast_detection(self, segmenter_vista_3d, test_images)` (line 213): Test contrast detection functionality. - - `def test_preprocessing(self, segmenter_vista_3d, test_images)` (line 241): Test preprocessing functionality. - - `def test_postprocessing(self, segmenter_vista_3d, test_images)` (line 260): Test postprocessing functionality. - - `def test_set_and_reset_prompts(self, segmenter_vista_3d)` (line 288): Test setting and resetting label prompt mode. - ## tests/test_segment_heart_simpleware.py - **class TestSegmentHeartSimpleware** (line 28): Test suite for SegmentHeartSimpleware (Simpleware Medical ASCardio). diff --git a/docs/README.md b/docs/README.md index 7028656..6f644d4 100644 --- a/docs/README.md +++ b/docs/README.md @@ -70,12 +70,10 @@ docs/ │ ├── index.rst # Main API hub │ ├── base.rst # Core base class │ ├── workflows.rst # Workflow classes -│ ├── segmentation/ # Segmentation (6 files) +│ ├── segmentation/ # Segmentation (4 files) │ │ ├── index.rst │ │ ├── base.rst │ │ ├── totalsegmentator.rst -│ │ ├── vista3d.rst -│ │ ├── vista3d_nim.rst │ │ └── ensemble.rst │ ├── registration/ # Image registration (5 files) │ │ ├── index.rst diff --git a/docs/api/index.rst b/docs/api/index.rst index 4ff8f53..c3ab4f6 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -20,8 +20,6 @@ This section provides detailed documentation for all PhysioMotion4D classes, fun segmentation/index segmentation/base segmentation/totalsegmentator - segmentation/vista3d - segmentation/vista3d_nim segmentation/ensemble .. toctree:: @@ -82,8 +80,7 @@ By Category **Segmentation** * :class:`~physiomotion4d.SegmentAnatomyBase` - Base segmentation class * :class:`~physiomotion4d.SegmentChestTotalSegmentator` - TotalSegmentator - * :class:`~physiomotion4d.SegmentChestVista3D` - VISTA-3D model - * :class:`~physiomotion4d.SegmentChestVista3DNIM` - VISTA-3D NIM + * :class:`~physiomotion4d.SegmentHeartSimpleware` - Simpleware cardiac segmentation * :class:`~physiomotion4d.SegmentChestEnsemble` - Ensemble segmentation **Image Registration** diff --git a/docs/api/segmentation/ensemble.rst b/docs/api/segmentation/ensemble.rst index a9eecde..620e8f1 100644 --- a/docs/api/segmentation/ensemble.rst +++ b/docs/api/segmentation/ensemble.rst @@ -4,7 +4,8 @@ Ensemble Segmentation .. currentmodule:: physiomotion4d -Combine multiple segmentation methods for improved accuracy and robustness. +Ensemble segmentation provides a stable API entry point that currently delegates +to :class:`SegmentChestTotalSegmentator`. Class Reference =============== @@ -18,59 +19,27 @@ Class Reference Overview ======== -Ensemble segmentation combines predictions from multiple methods using voting or averaging strategies to achieve higher accuracy than any single method. +``SegmentChestEnsemble`` inherits from :class:`SegmentChestTotalSegmentator` and +exposes the same interface. Using this class keeps downstream code stable if the +ensemble strategy changes in a future release. -**Key Features**: - * Combines TotalSegmentator and VISTA-3D - * Voting or averaging fusion strategies - * Improved boundary delineation - * Higher confidence in predictions - * Better robustness to image variations - -Usage Examples -============== - -Basic Ensemble --------------- +Usage Example +============= .. code-block:: python from physiomotion4d import SegmentChestEnsemble - - # Combine methods with voting - segmentator = SegmentChestEnsemble( - methods=['totalsegmentator', 'vista3d'], - fusion_strategy='voting', - verbose=True - ) - - labelmap = segmentator.segment("ct_scan.nrrd") - - # Get confidence map - confidence = segmentator.get_confidence_map() - -With Weighted Fusion --------------------- - -.. code-block:: python - # Weight methods differently - segmentator = SegmentChestEnsemble( - methods=['totalsegmentator', 'vista3d'], - fusion_strategy='weighted', - weights=[0.4, 0.6], # Trust VISTA-3D more - verbose=True - ) - - labelmap = segmentator.segment("cardiac_ct.nrrd") + segmenter = SegmentChestEnsemble() + result = segmenter.segment(ct_image, contrast_enhanced_study=False) + labelmap = result['labelmap'] See Also ======== * :doc:`index` - Segmentation overview -* :doc:`totalsegmentator` - Fast baseline method -* :doc:`vista3d` - High-accuracy method +* :doc:`totalsegmentator` - Underlying segmentation method .. rubric:: Navigation -:doc:`vista3d_nim` | :doc:`index` | :doc:`../registration/index` +:doc:`totalsegmentator` | :doc:`index` | :doc:`../registration/index` diff --git a/docs/api/segmentation/index.rst b/docs/api/segmentation/index.rst index d326ef1..2c97f2e 100644 --- a/docs/api/segmentation/index.rst +++ b/docs/api/segmentation/index.rst @@ -12,9 +12,8 @@ Overview PhysioMotion4D supports multiple segmentation approaches: * **TotalSegmentator**: Whole-body CT segmentation (100+ structures) -* **VISTA-3D**: MONAI-based foundation model for medical imaging -* **VISTA-3D NIM**: NVIDIA Inference Microservice version -* **Ensemble**: Combine multiple methods for improved accuracy +* **Simpleware**: Cardiac-focused segmentation (requires Simpleware Medical) +* **Ensemble**: Combines TotalSegmentator for improved robustness All segmentation classes inherit from :class:`SegmentAnatomyBase` and provide consistent interfaces. @@ -24,8 +23,6 @@ Quick Links **Segmentation Classes**: * :doc:`base` - Base class for all segmentation methods * :doc:`totalsegmentator` - TotalSegmentator implementation - * :doc:`vista3d` - VISTA-3D foundation model - * :doc:`vista3d_nim` - VISTA-3D NIM for cloud deployment * :doc:`ensemble` - Ensemble segmentation Choosing a Method @@ -36,9 +33,7 @@ Choosing a Method +==================+==================+==================+==================+ | TotalSegmentator | Fast (~30s) | Good | General purpose | +------------------+------------------+------------------+------------------+ -| VISTA-3D | Medium (~60s) | Excellent | Cardiac imaging | -+------------------+------------------+------------------+------------------+ -| VISTA-3D NIM | Fast (cloud) | Excellent | Production | +| Simpleware | Medium | Excellent | Cardiac imaging | +------------------+------------------+------------------+------------------+ | Ensemble | Slow (~90s) | Best | Research/QC | +------------------+------------------+------------------+------------------+ @@ -52,38 +47,10 @@ Basic Segmentation .. code-block:: python from physiomotion4d import SegmentChestTotalSegmentator - - # Initialize segmentator - segmentator = SegmentChestTotalSegmentator(fast=True, verbose=True) - - # Segment image - labelmap = segmentator.segment("ct_scan.nrrd") - - # Extract specific structure - heart = segmentator.extract_structure(labelmap, "heart") - - # Save results - segmentator.save_labelmap(labelmap, "output_labels.mha") - -With VISTA-3D -------------- - -.. code-block:: python - from physiomotion4d import SegmentChestVista3D - - # Initialize with GPU - segmentator = SegmentChestVista3D( - device="cuda:0", - use_auto_prompts=True, - verbose=True - ) - - # Segment specific structures - labelmap = segmentator.segment( - image_path="cardiac_ct.nrrd", - structures=["heart_left_ventricle", "heart_myocardium"] - ) + segmenter = SegmentChestTotalSegmentator() + result = segmenter.segment(ct_image, contrast_enhanced_study=False) + labelmap = result['labelmap'] Ensemble Approach ----------------- @@ -91,15 +58,10 @@ Ensemble Approach .. code-block:: python from physiomotion4d import SegmentChestEnsemble - - # Combine multiple methods - segmentator = SegmentChestEnsemble( - methods=['totalsegmentator', 'vista3d'], - fusion_strategy='voting', - verbose=True - ) - - labelmap = segmentator.segment("ct_scan.nrrd") + + segmenter = SegmentChestEnsemble() + result = segmenter.segment(ct_image, contrast_enhanced_study=False) + labelmap = result['labelmap'] Module Documentation ==================== @@ -109,8 +71,6 @@ Module Documentation base totalsegmentator - vista3d - vista3d_nim ensemble Common Operations @@ -123,17 +83,10 @@ Extract individual anatomical structures from segmentation results: .. code-block:: python - # Segment entire image - labelmap = segmentator.segment("ct.nrrd") - - # Extract cardiac structures - lv = segmentator.extract_structure(labelmap, "heart_left_ventricle") - rv = segmentator.extract_structure(labelmap, "heart_right_ventricle") - myocardium = segmentator.extract_structure(labelmap, "heart_myocardium") - - # Compute volumes - lv_volume = segmentator.compute_volume(lv) - print(f"LV volume: {lv_volume} mm³") + result = segmenter.segment(ct_image) + heart_mask = result['heart'] + lung_mask = result['lung'] + bone_mask = result['bone'] Batch Processing ---------------- @@ -143,53 +96,15 @@ Process multiple images efficiently: .. code-block:: python from pathlib import Path - - segmentator = SegmentChestTotalSegmentator(fast=True) - - for image_file in Path("data").glob("*.nrrd"): - print(f"Segmenting {image_file}...") - - labelmap = segmentator.segment(str(image_file)) - - output_file = f"{image_file.stem}_labels.mha" - segmentator.save_labelmap(labelmap, output_file) - -Quality Control ---------------- - -Validate segmentation quality: - -.. code-block:: python - - import numpy as np - - def validate_segmentation(labelmap, ground_truth): - """Compute Dice coefficient.""" - intersection = np.logical_and(labelmap, ground_truth).sum() - dice = 2 * intersection / (labelmap.sum() + ground_truth.sum()) - return dice - - # Validate results - dice_score = validate_segmentation(labelmap, reference_labelmap) - print(f"Dice score: {dice_score:.3f}") - -Best Practices -============== - -Method Selection ----------------- - -* **TotalSegmentator**: Use for general-purpose segmentation, fast iterations -* **VISTA-3D**: Use for cardiac structures, when accuracy is critical -* **Ensemble**: Use when maximum accuracy is needed, can tolerate longer processing + import itk -Parameter Tuning ----------------- + segmenter = SegmentChestTotalSegmentator() -* Start with default parameters -* Enable ``fast`` mode for quick prototyping -* Use GPU (``device="cuda:0"``) for VISTA-3D when available -* Adjust post-processing based on image quality + for image_file in Path("data").glob("*.nrrd"): + image = itk.imread(str(image_file)) + result = segmenter.segment(image) + labelmap = result['labelmap'] + itk.imwrite(labelmap, f"{image_file.stem}_labels.mha") Error Handling -------------- @@ -197,20 +112,9 @@ Error Handling .. code-block:: python try: - labelmap = segmentator.segment(image_path) + result = segmenter.segment(image) except RuntimeError as e: print(f"Segmentation failed: {e}") - # Fallback to alternative method - segmentator_backup = SegmentChestTotalSegmentator(fast=True) - labelmap = segmentator_backup.segment(image_path) - -Performance Tips -================ - -* Use GPU acceleration when available -* Enable fast mode for development -* Process time series in batch -* Cache segmentation results for repeated use See Also ======== diff --git a/docs/api/segmentation/vista3d.rst b/docs/api/segmentation/vista3d.rst deleted file mode 100644 index 791c081..0000000 --- a/docs/api/segmentation/vista3d.rst +++ /dev/null @@ -1,80 +0,0 @@ -==================================== -VISTA-3D Foundation Model -==================================== - -.. currentmodule:: physiomotion4d - -MONAI-based foundation model for high-accuracy medical image segmentation. - -Class Reference -=============== - -.. autoclass:: SegmentChestVista3D - :members: - :undoc-members: - :show-inheritance: - :inherited-members: - -Overview -======== - -VISTA-3D is a foundation model architecture that provides state-of-the-art accuracy, especially for cardiac structures. It supports interactive segmentation with point prompts and bounding boxes. - -**Key Features**: - * Foundation model trained on diverse medical datasets - * Interactive prompting with points and boxes - * Excellent accuracy on cardiac structures - * Supports automatic prompt generation - * Requires GPU for optimal performance - -Usage Examples -============== - -Basic Usage ------------ - -.. code-block:: python - - from physiomotion4d import SegmentChestVista3D - - # Initialize with GPU - segmentator = SegmentChestVista3D( - device="cuda:0", - verbose=True - ) - - # Segment with automatic prompts - labelmap = segmentator.segment( - image_path="cardiac_ct.nrrd", - structures=["heart_left_ventricle", "heart_myocardium"] - ) - -With Manual Prompts -------------------- - -.. code-block:: python - - # Segment with point prompts - labelmap = segmentator.segment( - image_path="ct.nrrd", - point_prompts=[(128, 128, 150)], # (x, y, z) coordinates - structure_name="heart_left_ventricle" - ) - - # Or with bounding box - labelmap = segmentator.segment( - image_path="ct.nrrd", - bbox_prompt=[100, 100, 120, 150, 150, 180], # xmin,ymin,zmin,xmax,ymax,zmax - structure_name="heart" - ) - -See Also -======== - -* :doc:`index` - Segmentation overview -* :doc:`vista3d_nim` - Cloud deployment version -* :doc:`totalsegmentator` - Alternative method - -.. rubric:: Navigation - -:doc:`totalsegmentator` | :doc:`index` | :doc:`vista3d_nim` diff --git a/docs/api/segmentation/vista3d_nim.rst b/docs/api/segmentation/vista3d_nim.rst deleted file mode 100644 index 61d720a..0000000 --- a/docs/api/segmentation/vista3d_nim.rst +++ /dev/null @@ -1,60 +0,0 @@ -========================================== -VISTA-3D NIM (Inference Microservice) -========================================== - -.. currentmodule:: physiomotion4d - -NVIDIA Inference Microservice version of VISTA-3D for cloud deployment. - -Class Reference -=============== - -.. autoclass:: SegmentChestVista3DNIM - :members: - :undoc-members: - :show-inheritance: - :inherited-members: - -Overview -======== - -VISTA-3D NIM provides optimized inference through a REST API, ideal for production deployments and scalable cloud applications. - -**Key Features**: - * Optimized for high-throughput inference - * REST API interface - * Cloud/server deployment - * Scalable for multiple concurrent requests - -Usage Examples -============== - -Basic Usage ------------ - -.. code-block:: python - - from physiomotion4d import SegmentChestVista3DNIM - - # Initialize with API endpoint - segmentator = SegmentChestVista3DNIM( - api_endpoint="https://api.nvidia.com/nim/vista3d", - api_key="your_api_key", - verbose=True - ) - - # Segment via API - labelmap = segmentator.segment( - image_path="ct_scan.nrrd", - structures=["heart", "lungs"] - ) - -See Also -======== - -* :doc:`vista3d` - Local deployment version -* :doc:`index` - Segmentation overview - -.. rubric:: Navigation - -:doc:`vista3d` | :doc:`index` | :doc:`ensemble` diff --git a/docs/developer/architecture.rst b/docs/developer/architecture.rst index 203b77d..e886aec 100644 --- a/docs/developer/architecture.rst +++ b/docs/developer/architecture.rst @@ -51,8 +51,6 @@ The package is organized into functional modules: ├── Segmentation │ ├── segment_anatomy_base.py Base segmentation │ ├── segment_chest_total_segmentator.py TotalSegmentator - │ ├── segment_chest_vista_3d.py VISTA-3D - │ ├── segment_chest_vista_3d_nim.py VISTA-3D NIM │ └── segment_chest_ensemble.py Ensemble methods │ ├── Registration @@ -100,7 +98,6 @@ Most PhysioMotion4D classes inherit from :class:`PhysioMotion4DBase`: ├── Segmentation Classes │ ├── SegmentAnatomyBase │ │ ├── SegmentChestTotalSegmentator - │ │ ├── SegmentChestVista3D │ │ └── SegmentChestEnsemble ├── Registration Classes │ ├── RegisterImagesBase @@ -275,7 +272,7 @@ Optional Dependencies * **ANTsPy**: ANTs registration (alternative to Icon) * **TotalSegmentator**: AI segmentation backend -* **MONAI**: Medical imaging AI framework (VISTA-3D) +* **MONAI**: Medical imaging AI framework * **CuPy**: GPU-accelerated array operations See ``pyproject.toml`` for complete dependency list. @@ -287,7 +284,7 @@ GPU Acceleration ---------------- * Registration (Icon): 5-10x speedup with GPU -* Segmentation (VISTA-3D): Requires GPU +* Segmentation: Requires GPU for best performance * Automatic GPU detection and fallback to CPU Memory Management diff --git a/docs/developer/segmentation.rst b/docs/developer/segmentation.rst index 656cd00..71f73a5 100644 --- a/docs/developer/segmentation.rst +++ b/docs/developer/segmentation.rst @@ -17,8 +17,6 @@ Overview PhysioMotion4D supports multiple segmentation approaches: * **TotalSegmentator**: Whole-body CT segmentation (100+ structures) -* **VISTA-3D**: MONAI-based foundation model for medical imaging -* **VISTA-3D NIM**: NVIDIA Inference Microservice version * **Ensemble**: Combine multiple methods for improved accuracy All segmentation classes inherit from :class:`SegmentAnatomyBase` and provide consistent interfaces. @@ -93,83 +91,6 @@ Uses the TotalSegmentator model for comprehensive anatomical segmentation. * Liver, kidneys, spleen * And many more... -VISTA-3D --------- - -MONAI-based foundation model for medical image segmentation. - -.. autoclass:: physiomotion4d.SegmentChestVista3D - :members: - :undoc-members: - :show-inheritance: - -**Features**: - * Foundation model architecture - * Supports point prompts and bounding boxes - * High accuracy on cardiac structures - * Requires GPU - -**Example Usage**: - -.. code-block:: python - - from physiomotion4d import SegmentChestVista3D - - # Initialize with GPU - segmentator = SegmentChestVista3D( - device="cuda:0", - use_auto_prompts=True, - verbose=True - ) - - # Segment with automatic prompts - labelmap = segmentator.segment( - image_path="cardiac_ct.nrrd", - structures=["heart_left_ventricle", "heart_myocardium"] - ) - - # Or provide manual prompts - labelmap = segmentator.segment( - image_path="cardiac_ct.nrrd", - point_prompts=[(128, 128, 150)], # (x, y, z) - structure_name="heart_left_ventricle" - ) - -VISTA-3D NIM ------------- - -NVIDIA Inference Microservice version for cloud deployment. - -.. autoclass:: physiomotion4d.SegmentChestVista3DNIM - :members: - :undoc-members: - :show-inheritance: - -**Features**: - * Optimized inference - * Cloud/server deployment - * REST API interface - * Scalable for multiple requests - -**Example Usage**: - -.. code-block:: python - - from physiomotion4d import SegmentChestVista3DNIM - - # Initialize with API endpoint - segmentator = SegmentChestVista3DNIM( - api_endpoint="https://api.nvidia.com/nim/vista3d", - api_key="your_api_key", - verbose=True - ) - - # Segment via API - labelmap = segmentator.segment( - image_path="ct_scan.nrrd", - structures=["heart", "lungs"] - ) - Ensemble Segmentation --------------------- @@ -181,7 +102,7 @@ Combines multiple segmentation methods for improved accuracy. :show-inheritance: **Features**: - * Combines TotalSegmentator + VISTA-3D + * Combines multiple segmentation methods * Voting or averaging strategies * Improved robustness * Better boundary delineation @@ -191,10 +112,10 @@ Combines multiple segmentation methods for improved accuracy. .. code-block:: python from physiomotion4d import SegmentChestEnsemble - + # Initialize ensemble segmentator = SegmentChestEnsemble( - methods=['totalsegmentator', 'vista3d'], + methods=['totalsegmentator'], fusion_strategy='voting', # or 'averaging' verbose=True ) @@ -359,7 +280,7 @@ Leverage GPU for faster inference: print("Using CPU") # Initialize with GPU - segmentator = SegmentChestVista3D(device=device) + segmentator = SegmentChestTotalSegmentator() Batch Processing ---------------- @@ -414,22 +335,6 @@ Assess segmentation quality: 'volume_difference': vol_diff } -Confidence Assessment ---------------------- - -Assess segmentation confidence: - -.. code-block:: python - - # For VISTA-3D (supports confidence) - segmentator = SegmentChestVista3D() - - labelmap, confidence_map = segmentator.segment_with_confidence("ct.nrrd") - - # Check low-confidence regions - low_confidence_mask = confidence_map < 0.5 - print(f"Low confidence regions: {low_confidence_mask.sum()} voxels") - Best Practices ============== @@ -437,7 +342,6 @@ Method Selection ---------------- * **TotalSegmentator**: General purpose, fast, comprehensive -* **VISTA-3D**: High accuracy, especially for cardiac structures * **Ensemble**: When accuracy is critical, can tolerate longer processing Parameter Tuning diff --git a/docs/examples.rst b/docs/examples.rst index 2fb4b6d..e078f37 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -93,28 +93,6 @@ Quick segmentation with TotalSegmentator: itk.imwrite(heart, "heart_mask.nrrd") itk.imwrite(lungs, "lungs_mask.nrrd") -VISTA-3D with Point Prompts ----------------------------- - -Advanced segmentation with user-provided points: - -.. code-block:: python - - from physiomotion4d import SegmentChestVista3D - import itk - - segmenter = SegmentChestVista3D() - image = itk.imread("chest_ct.nrrd") - - # Define point prompts (x, y, z in voxel coordinates) - heart_points = [(120, 150, 80), (130, 160, 85)] - - masks = segmenter.segment( - image, - contrast_enhanced_study=True, - point_prompts=heart_points - ) - Ensemble Segmentation --------------------- @@ -126,7 +104,7 @@ Combine multiple methods for best results: import itk segmenter = SegmentChestEnsemble( - methods=['totalsegmentator', 'vista3d'], + methods=['totalsegmentator'], fusion_strategy='voting' ) @@ -466,13 +444,13 @@ Segment multiple images in parallel: .. code-block:: python - from physiomotion4d import SegmentChestVista3D + from physiomotion4d import SegmentChestTotalSegmentator import itk import glob from concurrent.futures import ProcessPoolExecutor def segment_image(filename): - segmenter = SegmentChestVista3D() + segmenter = SegmentChestTotalSegmentator() image = itk.imread(filename) masks = segmenter.segment(image, contrast_enhanced_study=True) @@ -559,7 +537,7 @@ Mix and match different components: .. code-block:: python from physiomotion4d import ( - SegmentChestVista3D, + SegmentChestTotalSegmentator, RegisterImagesICON, TransformTools, ConvertVTKToUSDPolyMesh, @@ -572,7 +550,7 @@ Mix and match different components: frames = [itk.imread(f"frame_{i:03d}.mha") for i in range(10)] # Segment reference - segmenter = SegmentChestVista3D() + segmenter = SegmentChestTotalSegmentator() masks = segmenter.segment(reference, contrast_enhanced_study=True) heart_mask = masks[0] diff --git a/docs/index.rst b/docs/index.rst index 0de1edf..abd4c63 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -20,9 +20,9 @@ PhysioMotion4D is a comprehensive medical imaging package that converts 3D and 4 :target: https://github.com/Project-MONAI/physiomotion4d/blob/main/LICENSE :alt: License -.. image:: https://img.shields.io/github/actions/workflow/status/Project-MONAI/physiomotion4d/ci.yml?branch=main&label=CI%20Tests - :target: https://github.com/Project-MONAI/physiomotion4d/actions/workflows/ci.yml - :alt: CI Tests +.. image:: https://img.shields.io/github/actions/workflow/status/Project-MONAI/physiomotion4d/nightly-health.yml?branch=main&label=Nightly%20CI%20Tests + :target: https://github.com/Project-MONAI/physiomotion4d/actions/workflows/nightly-health.yml + :alt: Nightly CI Tests .. image:: https://img.shields.io/badge/tests-Windows%20%7C%20Linux%20%7C%20Python%203.10--3.12-blue :target: https://github.com/Project-MONAI/physiomotion4d/actions/workflows/ci.yml diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 9601ece..27809e3 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -116,11 +116,11 @@ If you only need segmentation: .. code-block:: python - from physiomotion4d import SegmentChestVista3D + from physiomotion4d import SegmentChestTotalSegmentator import itk # Initialize segmenter - segmenter = SegmentChestVista3D() + segmenter = SegmentChestTotalSegmentator() # Load and segment image image = itk.imread("chest_ct.nrrd") diff --git a/docs/testing.rst b/docs/testing.rst index 1e7ac4c..a6b4b79 100644 --- a/docs/testing.rst +++ b/docs/testing.rst @@ -48,7 +48,6 @@ PhysioMotion4D uses pytest markers to categorize tests: # Skip GPU-dependent tests pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py \ - --ignore=tests/test_segment_chest_vista_3d.py \ --ignore=tests/test_register_images_icon.py Specific Test Modules @@ -77,7 +76,6 @@ Specific Test Modules # Segmentation (GPU required, ~2-5 minutes each) pytest tests/test_segment_chest_total_segmentator.py -v - pytest tests/test_segment_chest_vista_3d.py -v Coverage Reports ---------------- @@ -103,8 +101,7 @@ Tests are organized by functionality: │ └── test_convert_nrrd_4d_to_3d.py # 4D to 3D conversion │ ├── Segmentation Tests (GPU Required) - │ ├── test_segment_chest_total_segmentator.py # TotalSegmentator - │ └── test_segment_chest_vista_3d.py # VISTA-3D segmentation + │ └── test_segment_chest_total_segmentator.py # TotalSegmentator │ ├── Registration Tests (Slow ~5-10 min) │ ├── test_register_images_ants.py # ANTs registration diff --git a/experiments/Heart-GatedCT_To_USD/test_vista3d_class.py b/experiments/Heart-GatedCT_To_USD/test_vista3d_class.py deleted file mode 100644 index 9147d51..0000000 --- a/experiments/Heart-GatedCT_To_USD/test_vista3d_class.py +++ /dev/null @@ -1,22 +0,0 @@ -#!/usr/bin/env python -# %% -import os - -import itk - -from physiomotion4d.segment_chest_vista_3d import SegmentChestVista3D - -_HERE = os.path.dirname(os.path.abspath(__file__)) - -output_dir = os.path.join(_HERE, "results") -max_image = itk.imread(os.path.join(output_dir, "slice_fixed.mha")) - -# %% -seg = SegmentChestVista3D() -result = seg.segment(max_image, contrast_enhanced_study=True) -labelmap_image = result["labelmap"] -itk.imwrite( - labelmap_image, - os.path.join(output_dir, "slice_fixed.all_mask_vista3d.mha"), - compression=True, -) diff --git a/experiments/Heart-GatedCT_To_USD/test_vista3d_inMem.py b/experiments/Heart-GatedCT_To_USD/test_vista3d_inMem.py deleted file mode 100644 index 4c9c49c..0000000 --- a/experiments/Heart-GatedCT_To_USD/test_vista3d_inMem.py +++ /dev/null @@ -1,207 +0,0 @@ -#!/usr/bin/env python -# %% -import itk -import numpy as np -import torch - - -def vista3d_inference_from_itk( - itk_image, - label_prompt=None, - points=None, - point_labels=None, - device=None, - bundle_path=None, - model_cache_dir=None, -): - # 1. Import dependencies - import itk - from monai.bundle import download - from monai.data.itk_torch_bridge import itk_image_to_metatensor - from monai.inferers import sliding_window_inference - from monai.networks.nets import vista3d132 - from monai.transforms import ( - CropForeground, - EnsureChannelFirst, - EnsureType, - ScaleIntensityRange, - Spacing, - ) - from monai.utils import set_determinism - - set_determinism(seed=42) - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - - # 2. Handle "no prompts" case: segment all classes - if label_prompt is None and points is None: - everything_labels = list( - set([i + 1 for i in range(132)]) - set([2, 16, 18, 20, 21, 23, 24, 25, 26]) - ) - label_prompt = everything_labels - print( - f"No prompt provided. Using everything_labels for {len(everything_labels)} classes." - ) - - if points is not None and point_labels is None: - raise ValueError("point_labels must be provided when points are specified") - - # 3. Download model bundle if needed - if bundle_path is None: - import tempfile - - if model_cache_dir is None: - model_cache_dir = tempfile.mkdtemp() - try: - download(name="vista3d", bundle_dir=model_cache_dir, source="monaihosting") - except Exception: - download(name="vista3d", bundle_dir=model_cache_dir, source="github") - bundle_path = f"{model_cache_dir}/vista3d" - - # 4. ITK->MetaTensor (in memory) - meta_tensor = itk_image_to_metatensor( - itk_image, channel_dim=None, dtype=torch.float32 - ) - - # 5. Preprocessing pipeline - processed = meta_tensor - processed = EnsureChannelFirst(channel_dim=None)(processed) - processed = EnsureType(dtype=torch.float32)(processed) - processed = Spacing(pixdim=[1.5, 1.5, 1.5], mode="bilinear")(processed) - processed = ScaleIntensityRange( - a_min=-1024, a_max=1024, b_min=0.0, b_max=1.0, clip=True - )(processed) - processed = CropForeground()(processed) - - # Save the MONAI affine now (Spacing + CropForeground have updated it). - # We need it later to place the label map in the processed world space before - # resampling back to the original ITK grid. - processed_affine = ( - processed.meta["affine"].numpy() - if hasattr(processed, "meta") and "affine" in processed.meta - else None - ) - - # 6. Load VISTA3D - model = vista3d132(encoder_embed_dim=48, in_channels=1) - model_path = f"{bundle_path}/models/model.pt" - checkpoint = torch.load(model_path, map_location=device) - model.load_state_dict(checkpoint) - model.eval() - model.to(device) - - # 7. Prepare input tensor - input_tensor = processed - if not isinstance(input_tensor, torch.Tensor): - input_tensor = torch.tensor(np.asarray(input_tensor), dtype=torch.float32) - if input_tensor.dim() == 3: - input_tensor = input_tensor.unsqueeze(0) - if input_tensor.dim() == 4: - input_tensor = input_tensor.unsqueeze(0) - input_tensor = input_tensor.to(device) - - # 8. Prepare model inputs - model_inputs = {"image": input_tensor} - if label_prompt is not None: - label_prompt_tensor = torch.tensor( - label_prompt, dtype=torch.long, device=device - ) - model_inputs["label_prompt"] = label_prompt_tensor - print("label_prompt_tensor shape", label_prompt_tensor.shape) - if points is not None: - point_coords = torch.tensor( - points, dtype=torch.float32, device=device - ).unsqueeze(0) - point_labels_tensor = torch.tensor( - point_labels, dtype=torch.float32, device=device - ).unsqueeze(0) - model_inputs["points"] = point_coords - model_inputs["point_labels"] = point_labels_tensor - print("point_coords shape", point_coords.shape) - - # 9. Sliding window inference for large images - def predictor_fn(x): - args = {k: v for k, v in model_inputs.items() if k != "image"} - print(x.shape) - return model(x, **args) - - with torch.no_grad(): - if any(dim > 128 for dim in input_tensor.shape[2:]): - print("Sliding window inference") - output = sliding_window_inference( - input_tensor, - roi_size=[128, 128, 128], - sw_batch_size=1, - predictor=predictor_fn, - overlap=0.5, - mode="gaussian", - device=device, - ) - else: - print("Single window inference") - output = model( - input_tensor, **{k: v for k, v in model_inputs.items() if k != "image"} - ) - - print("output shape", output.shape) - # 10. Postprocess: multi-class to label map - output = output.cpu() - if hasattr(output, "detach"): - output = output.detach() - if isinstance(output, dict): - if "pred" in output: - output = output["pred"] - else: - output = list(output.values())[0] - - if output.shape[1] > 1: - label_map = torch.argmax(output, dim=1).squeeze(0).numpy().astype(np.uint16) - else: - label_map = (output > 0.5).squeeze(0).cpu().numpy().astype(np.uint8) - - # MONAI outputs are in (D, H, W) = (z, y, x) — matches ITK's GetImageFromArray - # convention, so no transpose is needed. - label_map_for_itk = label_map - - # Build an ITK image in the processed (1.5 mm, cropped) world space. - output_itk = itk.GetImageFromArray(label_map_for_itk) - if processed_affine is not None: - # Extract spacing and origin from the MONAI affine matrix. - # Columns norms of the 3×3 rotation-scale block give voxel spacing. - spacing_processed = np.sqrt( - (processed_affine[:3, :3] ** 2).sum(axis=0) - ).tolist() - origin_processed = processed_affine[:3, 3].tolist() - output_itk.SetSpacing(spacing_processed) - output_itk.SetOrigin(origin_processed) - output_itk.SetDirection(itk_image.GetDirection()) - - # Resample the label map back to the original input image grid using - # nearest-neighbour interpolation (preserves discrete label values). - resampler = itk.ResampleImageFilter.New(output_itk) - resampler.SetReferenceImage(itk_image) - resampler.SetUseReferenceImage(True) - resampler.SetInterpolator( - itk.NearestNeighborInterpolateImageFunction.New(output_itk) - ) - resampler.SetDefaultPixelValue(0) - resampler.Update() - output_itk = resampler.GetOutput() - else: - # Fallback: copy input metadata (spatial alignment may be approximate). - output_itk.SetSpacing(itk_image.GetSpacing()) - output_itk.SetOrigin(itk_image.GetOrigin()) - output_itk.SetDirection(itk_image.GetDirection()) - - return output_itk - - -# %% -# Load an ITK image -image = itk.imread("results/slice_fixed.mha") - -spleen_segmentation = vista3d_inference_from_itk( - image, model_cache_dir="./network_weights" -) - -itk.imwrite(spleen_segmentation, "results/slice_fixed.all_mask_vista3d_inMem.mha") diff --git a/experiments/Lung-GatedCT_To_USD/0-register_dirlab_4dct.py b/experiments/Lung-GatedCT_To_USD/0-register_dirlab_4dct.py index 5b4056c..2a08ad7 100644 --- a/experiments/Lung-GatedCT_To_USD/0-register_dirlab_4dct.py +++ b/experiments/Lung-GatedCT_To_USD/0-register_dirlab_4dct.py @@ -95,7 +95,6 @@ def register_image( # %% -# seg_image = SegmentChestVista3D() seg_image = SegmentChestTotalSegmentator() os.makedirs(output_dir, exist_ok=True) diff --git a/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py b/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py index cebb9e5..a428df7 100644 --- a/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py +++ b/experiments/Lung-GatedCT_To_USD/Experiment_SegReg.py @@ -8,7 +8,6 @@ from physiomotion4d import RegisterImagesICON from physiomotion4d import SegmentChestTotalSegmentator -from physiomotion4d import SegmentChestVista3D _HERE = os.path.dirname(os.path.abspath(__file__)) _DATA_DIR = os.path.join(_HERE, "..", "..", "data", "DirLab-4DCT") @@ -44,50 +43,3 @@ os.path.join(_RESULTS_DIR, "Experiment_totseg.mha"), compression=True, ) - -# %% -# This section requires the Vista3D container to be running - -vista3d_running = False -if vista3d_running: - img = itk.imread(os.path.join(_RESULTS_DIR, "Experiment_reg.mha")) - - tot_seg = SegmentChestVista3D() - - seg_image = tot_seg.segment(img, contrast_enhanced_study=False) - - itk.imwrite( - seg_image[0], - os.path.join(_RESULTS_DIR, "Experiment_vista3d.mha"), - compression=True, - ) - itk.imwrite( - seg_image[1], - os.path.join(_RESULTS_DIR, "Experiment_vista3d_lung.mha"), - compression=True, - ) - itk.imwrite( - seg_image[2], - os.path.join(_RESULTS_DIR, "Experiment_vista3d_heart.mha"), - compression=True, - ) - itk.imwrite( - seg_image[3], - os.path.join(_RESULTS_DIR, "Experiment_vista3d_bone.mha"), - compression=True, - ) - itk.imwrite( - seg_image[4], - os.path.join(_RESULTS_DIR, "Experiment_vista3d_soft_tissue.mha"), - compression=True, - ) - itk.imwrite( - seg_image[5], - os.path.join(_RESULTS_DIR, "Experiment_vista3d_other.mha"), - compression=True, - ) - itk.imwrite( - seg_image[6], - os.path.join(_RESULTS_DIR, "Experiment_vista3d_contrast.mha"), - compression=True, - ) diff --git a/experiments/README.md b/experiments/README.md index 275bf89..ebf5e64 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -64,7 +64,7 @@ visualization and manipulation. **Pipeline stages:** 1. 4D CT reconstruction (using methods from `Reconstruct4DCT`) -2. AI segmentation (TotalSegmentator, Vista3D/Clara Segment Open Model as NIM or local) +2. AI segmentation (TotalSegmentator) 3. Per-organ mesh generation 4. OpenUSD conversion with animation 5. Tissue property mapping (subsurface scatter, color, etc.) @@ -317,7 +317,7 @@ Each subdirectory represents a different experimental domain: This experimental code was instrumental in: 1. Defining the final library architecture 2. Testing registration algorithms (ICON, SyN, LDDMM) -3. Evaluating segmentation approaches (TotalSegmentator, VISTA-3D) +3. Evaluating segmentation approaches (TotalSegmentator) 4. Developing the USD export pipeline 5. Optimizing the complete 4D CT → USD workflow 6. Identifying modular extension points for new anatomical regions and tasks diff --git a/pyproject.toml b/pyproject.toml index e117666..28a28b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,6 @@ keywords = [ "ai", "deep-learning", "totalsegmentator", - "vista-3d", "icon", "ants", "physiological-motion", @@ -230,7 +229,6 @@ warn_no_return = true warn_unreachable = true strict_equality = true show_error_codes = true -exclude = '(?x)(^src[/\\\\]physiomotion4d[/\\\\]network_weights[/\\\\]vista3d[/\\\\])' # Third-party libs (itk, pyvista, pxr, vtk, simpleware, etc.) have no stubs disable_error_code = ["import-not-found", "import-untyped"] diff --git a/src/physiomotion4d/__init__.py b/src/physiomotion4d/__init__.py index 9f86afb..e139911 100644 --- a/src/physiomotion4d/__init__.py +++ b/src/physiomotion4d/__init__.py @@ -64,8 +64,6 @@ from .segment_anatomy_base import SegmentAnatomyBase from .segment_chest_ensemble import SegmentChestEnsemble from .segment_chest_total_segmentator import SegmentChestTotalSegmentator -from .segment_chest_vista_3d import SegmentChestVista3D -from .segment_chest_vista_3d_nim import SegmentChestVista3DNIM from .segment_heart_simpleware import SegmentHeartSimpleware from .transform_tools import TransformTools from .usd_anatomy_tools import USDAnatomyTools @@ -93,8 +91,6 @@ "SegmentAnatomyBase", "SegmentChestEnsemble", "SegmentChestTotalSegmentator", - "SegmentChestVista3D", - "SegmentChestVista3DNIM", "SegmentHeartSimpleware", # Registration classes "RegisterImagesBase", diff --git a/src/physiomotion4d/cli/convert_ct_to_vtk.py b/src/physiomotion4d/cli/convert_ct_to_vtk.py index 6d72cfa..fb7d4c7 100644 --- a/src/physiomotion4d/cli/convert_ct_to_vtk.py +++ b/src/physiomotion4d/cli/convert_ct_to_vtk.py @@ -43,15 +43,6 @@ def main() -> int: --input-image chest_ct.nii.gz \\ --output-dir ./results - # VISTA-3D, contrast-enhanced, split per group - %(prog)s \\ - --input-image chest_ct.nii.gz \\ - --segmentation-method vista_3d \\ - --contrast \\ - --split-files \\ - --output-dir ./results \\ - --output-prefix patient01 - # Simpleware heart-only, cardiac anatomy groups, combined output %(prog)s \\ --input-image chest_ct.nii.gz \\ @@ -85,10 +76,7 @@ def main() -> int: "--segmentation-method", default="total_segmentator", choices=list(WorkflowConvertCTToVTK.SEGMENTATION_METHODS), - help=( - "Segmentation backend. " - "total_segmentator (default) | vista_3d | simpleware_heart" - ), + help=("Segmentation backend. total_segmentator (default) | simpleware_heart"), ) parser.add_argument( "--contrast", diff --git a/src/physiomotion4d/segment_chest_ensemble.py b/src/physiomotion4d/segment_chest_ensemble.py index b51bb3c..a19ce45 100644 --- a/src/physiomotion4d/segment_chest_ensemble.py +++ b/src/physiomotion4d/segment_chest_ensemble.py @@ -1,418 +1,30 @@ -"""Module for segmenting chest CT images using VISTA3D.""" +"""Module for ensemble chest CT segmentation. -# Please start vista3d docker: -# docker run --rm -it --name vista3d --runtime=nvidia -# -e CUDA_VISIBLE_DEVICES=0 -# -e NGC_API_KEY=$NGC_API_KEY -# --shm-size=8G -p 8000:8000 -# -v /tmp/data:/home/aylward/tmp/data nvcr.io/nim/nvidia/vista3d:latest +Currently delegates to TotalSegmentator. The ensemble class is retained as a +public API entry point so that callers do not need to change their imports. +""" import logging -import itk -import numpy as np - -from physiomotion4d.segment_anatomy_base import SegmentAnatomyBase from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator -from physiomotion4d.segment_chest_vista_3d import SegmentChestVista3D - - -class SegmentChestEnsemble(SegmentAnatomyBase): - """ - A class that inherits from physioSegmentChest and implements the - segmentation method using VISTA3D. - """ - - def __init__(self, log_level: int | str = logging.INFO): - """Initialize the vista3d class. - - Args: - log_level: Logging level (default: logging.INFO) - """ - super().__init__(log_level=log_level) - - self.target_spacing = 0.0 - - self.heart_mask_ids = { - 108: "left_atrial_appendage", - 115: "heart", - 140: "heart_envelope", - } - - self.major_vessels_mask_ids = { - 6: "aorta", - 7: "inferior_vena_cava", - 17: "portal_vein_and_splenic_vein", - 58: "left_iliac_artery", - 59: "right_iliac_artery", - 60: "left_iliac_vena", - 61: "right_iliac_vena", - 110: "left_brachiocephalic_vena", - 111: "right_brachiocephalic_vena", - 112: "left_common_carotid_artery", - 113: "right_common_carotid_artery", - 119: "pulmonary_vein", - 123: "left_subclavian_artery", - 124: "right_subclavian_artery", - 125: "superior_vena_cava", - 146: "brachiocephalic_trunk", - } - - self.lung_mask_ids = { - 28: "left_lung_upper_lobe", - 29: "left_lung_lower_lobe", - 30: "right_lung_upper_lobe", - 31: "right_lung_middle_lobe", - 32: "right_lung_lower_lobe", - 57: "trachea", - 132: "airway", - 147: "esophagus", - } - - self.bone_mask_ids = { - 33: "vertebrae_L5", - 34: "vertebrae_L4", - 35: "vertebrae_L3", - 36: "vertebrae_L2", - 37: "vertebrae_L1", - 38: "vertebrae_T12", - 39: "vertebrae_T11", - 40: "vertebrae_T10", - 41: "vertebrae_T9", - 42: "vertebrae_T8", - 43: "vertebrae_T7", - 44: "vertebrae_T6", - 45: "vertebrae_T5", - 46: "vertebrae_T4", - 47: "vertebrae_T3", - 48: "vertebrae_T2", - 49: "vertebrae_T1", - 50: "vertebrae_C7", - 51: "vertebrae_C6", - 52: "vertebrae_C5", - 53: "vertebrae_C4", - 54: "vertebrae_C3", - 55: "vertebrae_C2", - 56: "vertebrae_C1", - 63: "left_rib_1", - 64: "left_rib_2", - 65: "left_rib_3", - 66: "left_rib_4", - 67: "left_rib_5", - 68: "left_rib_6", - 69: "left_rib_7", - 70: "left_rib_8", - 71: "left_rib_9", - 72: "left_rib_10", - 73: "left_rib_11", - 74: "left_rib_12", - 75: "right_rib_1", - 76: "right_rib_2", - 77: "right_rib_3", - 78: "right_rib_4", - 79: "right_rib_5", - 80: "right_rib_6", - 81: "right_rib_7", - 82: "right_rib_8", - 83: "right_rib_9", - 84: "right_rib_10", - 85: "right_rib_11", - 86: "right_rib_12", - 87: "left_humerus", - 88: "right_humerus", - 89: "left_scapula", - 90: "right_scapula", - 91: "left_clavicula", - 92: "right_clavicula", - 93: "left_femur", - 94: "right_femur", - 95: "left_hip", - 96: "right_hip", - 120: "skull", - 122: "sternum", - 114: "costal_cartilages", - 127: "vertebrae_S1", - } - - self.soft_tissue_mask_ids = { - 3: "spleen", - 5: "right_kidney", - 14: "left_kidney", - 10: "gallbladder", - 1: "liver", - 12: "stomach", - 4: "pancreas", - 8: "right_adrenal_gland", - 9: "left_adrenal_gland", - 126: "thyroid_gland", - 19: "small_bowel", - 13: "duodenum", - 62: "colon", - 15: "bladder", - 118: "prostate", - 121: "spinal_cord", - 22: "brain", - 133: "soft_tissue", - 148: "sacrum", - 149: "gluteus_maximus_left", - 150: "gluteus_maximus_right", - 151: "gluteus_medius_left", - 152: "gluteus_medius_right", - 153: "gluteus_minimus_left", - 154: "gluteus_minimus_right", - } - - self.set_other_and_all_mask_ids() - - self.heart_ids_map = { - 108: 61, # left_atrial_appendage / atrial_appendage_left - 115: 51, # heart - 140: 140, # heart_envelope / heart_envelop - } - - self.major_vessels_ids_map = { - 6: 52, # aorta - 7: 63, # inferior_vena_cava - 17: -1, # portal_vein_and_splenic_vein - 58: -1, # left_iliac_artery - 59: -1, # right_iliac_artery - 60: -1, # left_iliac_vena - 61: -1, # right_iliac_vena - 110: 59, # left_brachiocephalic_vena / brachiocephalic_vein_left - 111: 60, # right_brachiocephalic_vena / brachiocephalic_vein_right - 112: 58, # left_common_carotid_artery / common_carotid_artery_left - 113: 57, # right_common_carotid_artery / common_carotid_artery_right - 119: 53, # pulmonary_vein - 123: 56, # left_subclavian_artery - 124: 55, # right_subclavian_artery - 125: 62, # superior_vena_cava - 146: 54, # brachiocephalic_trunk (new key for only-in-2nd-list) - } - - self.lung_ids_map = { - 28: 10, # left_lung_upper_lobe / lung_upper_lobe_left - 29: 11, # left_lung_lower_lobe / lung_lower_lobe_left - 30: 12, # right_lung_upper_lobe / lung_upper_lobe_right - 31: 13, # right_lung_middle_lobe / lung_middle_lobe_right - 32: 14, # right_lung_lower_lobe / lung_lower_lobe_right - 57: 16, # trachea - 132: -1, # airway (only in list1) - 147: 15, # esophagus (only in list2) - } - self.bone_ids_map = { - 33: 27, # vertebrae L5 / vertebra_L5 - 34: 28, - 35: 29, - 36: 30, - 37: 31, - 38: 32, - 39: 33, - 40: 34, - 41: 35, - 42: 36, - 43: 37, - 44: 38, - 45: 39, - 46: 40, - 47: 41, - 48: 42, - 49: 43, - 50: 44, - 51: 45, - 52: 46, - 53: 47, - 54: 48, - 55: 49, - 56: 50, - 63: 92, # left_rib_1 / rib_left_1 - 64: 93, - 65: 94, - 66: 95, - 67: 96, - 68: 97, - 69: 98, - 70: 99, - 71: 100, - 72: 101, - 73: 102, - 74: 103, - 75: 104, # right_rib_1 / rib_right_1 - 76: 105, - 77: 106, - 78: 107, - 79: 108, - 80: 109, - 81: 110, - 82: 111, - 83: 112, - 84: 113, - 85: 114, - 86: 115, - 87: 69, # left_humerus / humerus_left - 88: 70, # right_humerus / humerus_right - 89: 71, # left_scapula / scapula_left - 90: 72, # right_scapula / scapula_right - 91: 73, # left_clavicula / clavicula_left - 92: 74, # right_clavicula / clavicula_right - 93: 75, # left_femur / femur_left - 94: 76, # right_femur / femur_right - 95: 77, # left_hip / hip_left - 96: 78, # right_hip / hip_right - 120: 91, # skull - 122: 116, # sternum - 114: 117, # costal_cartilages - 127: 26, # vertebrae_S1 / vertebra_S1 - } - self.soft_tissue_ids_map = { - 3: 1, # spleen - 5: 2, # right_kidney / kidney_right - 14: 3, # left_kidney / kidney_left - 10: 4, # gallbladder - 1: 5, # liver - 12: 6, # stomach - 4: 7, # pancreas - 8: 8, # right_adrenal_gland / adrenal_gland_right - 9: 9, # left_adrenal_gland / adrenal_gland_left - 126: 17, # thyroid_gland - 19: 18, # small_bowel - 13: 19, # duodenum - 62: 20, # colon - 15: 21, # bladder / urinary_bladder - 118: 22, # prostate - 121: -1, # spinal_cord (only in list1) - 22: 90, # brain - 133: 133, # soft_tissue - # Only-in-second-list below: - 148: 25, # sacrum (new unique key) - 149: 80, # gluteus_maximus_left (new unique key) - 150: 81, # gluteus_maximus_right (new unique key) - 151: 82, # gluteus_medius_left (new unique key) - 152: 83, # gluteus_medius_right (new unique key) - 153: 84, # gluteus_minimus_left (new unique key) - 154: 85, # gluteus_minimus_right (new unique key) - } +class SegmentChestEnsemble(SegmentChestTotalSegmentator): + """Ensemble chest CT segmentation. - self.vista3d_to_totseg_ids_map = { - **self.heart_ids_map, - **self.major_vessels_ids_map, - **self.lung_ids_map, - **self.bone_ids_map, - **self.soft_tissue_ids_map, - } + Inherits from :class:`SegmentChestTotalSegmentator` and currently delegates + all segmentation to that backend. The class is kept as a stable public + name so that downstream code depending on ``SegmentChestEnsemble`` continues + to work without modification. - self.totseg_to_vista3d_ids_map = { - v: k for k, v in self.vista3d_to_totseg_ids_map.items() if v != -1 - } + Args: + log_level: Logging level (default: ``logging.INFO``). - def ensemble_segmentation( - self, labelmap_vista: itk.image, labelmap_totseg: itk.image - ) -> itk.image: - """ - Combine two segmentation results using label mapping and priority rules. - - Args: - labelmap_vista (itk.image): The VISTA3D segmentation result. - labelmap_totseg (itk.image): The TotalSegmentator segmentation result. - - Returns: - itk.image: The combined segmentation result. - """ - - self.log_info("Running ensemble segmentation: combining results") - - labelmap_vista_arr = itk.GetArrayFromImage(labelmap_vista) - labelmap_totseg_arr = itk.GetArrayFromImage(labelmap_totseg) - self.log_info("Segmentations loaded") - - results_arr = np.zeros_like(labelmap_vista_arr) - - self.log_info("Setting interpolators") - labelmap_vista_interp = itk.LabelImageGaussianInterpolateImageFunction.New( - labelmap_vista - ) - labelmap_totseg_interp = itk.LabelImageGaussianInterpolateImageFunction.New( - labelmap_totseg - ) - - self.log_info("Iterating through labelmaps") - lastidx0 = -1 - total_slices = labelmap_vista_arr.shape[0] - for idx in np.ndindex(labelmap_vista_arr.shape): - if idx[0] != lastidx0: - if idx[0] % 10 == 0 or idx[0] == total_slices - 1: - self.log_progress( - idx[0] + 1, total_slices, prefix="Processing slices" - ) - lastidx0 = idx[0] - # Skip if both are zero - vista_label = labelmap_vista_arr[idx] - totseg_label = labelmap_totseg_arr[idx] - if vista_label == 0 and totseg_label == 0: - continue - - totseg_vista_label = self.totseg_to_vista3d_ids_map.get(totseg_label, 0) - if vista_label == 0: - results_arr[idx] = totseg_vista_label - elif totseg_label == 0: - results_arr[idx] = vista_label - else: - # print("Conflict detected at", idx, vista_label, totseg_label, end="", flush=True) - # Softtissue label in Vista3D is a catch-all label, - # so use the TotalSegmentator label instead - label = totseg_vista_label - if vista_label != 133: - for sigma in [2, 4, 8]: - labelmap_vista_interp.SetSigma([sigma, sigma, sigma]) - labelmap_totseg_interp.SetSigma([sigma, sigma, sigma]) - tmp_vista = labelmap_vista_interp.EvaluateAtIndex(idx[::-1]) - tmp_totseg = labelmap_totseg_interp.EvaluateAtIndex(idx[::-1]) - # print(" ", tmp_vista, tmp_totseg, end="", flush=True) - if tmp_vista == 0 and tmp_totseg == 0: - label = 0 - # print("...agreeing on 0", flush=True) - break - tmp_totseg_vista = self.totseg_to_vista3d_ids_map.get( - tmp_totseg, 0 - ) - # print(f"({tmp_totseg_vista})", end="", flush=True) - if tmp_vista == tmp_totseg_vista: - label = tmp_vista - # print("...agreeing on", label, flush=True) - break - label = totseg_vista_label - # print(" assigning =", label, flush=True) - results_arr[idx] = label - labelmap_vista_arr[idx] = label - totseg_label = self.vista3d_to_totseg_ids_map.get(vista_label, 0) - totseg_label = max(totseg_label, 0) - labelmap_totseg_arr[idx] = totseg_label - - results_arr = results_arr.reshape(labelmap_vista_arr.shape) - results_image = itk.GetImageFromArray(results_arr) - results_image.CopyInformation(labelmap_vista) - - return results_image - - def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: - """ - Run VISTA3D on the preprocessed image and return result. - - Args: - preprocessed_image (itk.image): The preprocessed image to segment. - - Returns: - the segmented image. - """ - vista3d_result = SegmentChestVista3D().segment(preprocessed_image) - vista3d_labelmap = vista3d_result["labelmap"] - - total_result = SegmentChestTotalSegmentator().segment(preprocessed_image) - total_labelmap = total_result["labelmap"] - - ensemble_segmentation = self.ensemble_segmentation( - vista3d_labelmap, total_labelmap - ) + Example: + >>> segmenter = SegmentChestEnsemble() + >>> result = segmenter.segment(ct_image, contrast_enhanced_study=False) + >>> labelmap = result['labelmap'] + """ - return ensemble_segmentation + def __init__(self, log_level: int | str = logging.INFO) -> None: + super().__init__(log_level=log_level) diff --git a/src/physiomotion4d/segment_chest_vista_3d.py b/src/physiomotion4d/segment_chest_vista_3d.py deleted file mode 100644 index 2373178..0000000 --- a/src/physiomotion4d/segment_chest_vista_3d.py +++ /dev/null @@ -1,411 +0,0 @@ -"""Module for segmenting chest CT images using VISTA3D. - -This module provides the SegmentChestVista3D class that implements chest CT -segmentation using the VISTA-3D foundational model from NVIDIA. VISTA-3D is -a versatile segmentation model that can perform both automatic segmentation -and interactive segmentation with point or label prompts. - -The module requires the VISTA-3D model weights to be downloaded from Hugging Face -and supports both local inference and NVIDIA NIM deployment modes. -""" - -# Please start vista3d docker: -# docker run --rm -it --name vista3d --runtime=nvidia -# -e CUDA_VISIBLE_DEVICES=0 -# -e NGC_API_KEY=$NGC_API_KEY -# --shm-size=8G -p 8000:8000 -# -v /tmp/data:/home/aylward/tmp/data nvcr.io/nim/nvidia/vista3d:latest - -import logging -import os -import sys -import tempfile -from typing import Optional - -import itk -import torch -from huggingface_hub import snapshot_download - -from physiomotion4d.segment_anatomy_base import SegmentAnatomyBase - - -class SegmentChestVista3D(SegmentAnatomyBase): - """ - Chest CT segmentation using NVIDIA VISTA-3D foundational model. - - This class implements chest CT segmentation using the VISTA-3D model, - a versatile foundational segmentation model that supports both automatic - ('everything') segmentation and interactive segmentation with prompts. - - VISTA-3D is a state-of-the-art 3D medical image segmentation model that - can segment 132+ anatomical structures. It supports two interaction modes: - 1. Everything segmentation: Segments all detectable structures - 2. Label prompts: Segments specific structures by ID - - The class automatically downloads model weights from Hugging Face and - supports GPU acceleration. It includes additional soft tissue segmentation - to fill gaps not covered by the base VISTA-3D segmentation. - - Attributes: - target_spacing (float): Adaptive spacing based on input image - device (torch.device): GPU device for model inference - bundle_path (str): Path to VISTA-3D model weights - hf_pipeline (object): Hugging Face pipeline for inference - label_prompt (list): Specific anatomical structure IDs to segment - - Example: - >>> # Automatic segmentation - >>> segmenter = SegmentChestVista3D() - >>> result = segmenter.segment(ct_image, contrast_enhanced_study=True) - >>> labelmap = result['labelmap'] - >>> heart_mask = result['heart'] - >>> - >>> # Segment specific structures - >>> segmenter.set_label_prompt([115, 6, 28]) # Heart, aorta, left lung - >>> result = segmenter.segment(ct_image) - """ - - def __init__(self, log_level: int | str = logging.INFO) -> None: - """Initialize the VISTA-3D based chest segmentation. - - Sets up the VISTA-3D model including downloading weights from Hugging Face, - configuring GPU device, and initializing anatomical structure mappings - specific to VISTA-3D's label set. - - The initialization automatically downloads the VISTA-3D model weights - from the MONAI/VISTA3D-HF repository on Hugging Face if not already present. - - Args: - log_level: Logging level (default: logging.INFO) - - Raises: - RuntimeError: If CUDA is not available for GPU acceleration - ConnectionError: If model weights cannot be downloaded - """ - super().__init__(log_level=log_level) - - self.target_spacing = 0.0 - self.resale_intensity_range = False - self.input_percentile_range = None - self.output_percentile_range = None - self.output_intensity_range = None - - self.device = torch.device("cuda:0") - - self.bundle_path = os.path.join( - os.path.dirname(__file__), "network_weights/vista3d" - ) - os.makedirs(self.bundle_path, exist_ok=True) - - self.model_name = "vista3d" - - self.label_prompt: Optional[list[int]] = None - - repo_id = "MONAI/VISTA3D-HF" - snapshot_download(repo_id=repo_id, local_dir=self.bundle_path) - - self.heart_mask_ids = { - 108: "left_atrial_appendage", - 115: "heart", - 140: "heart_envelope", - } - - self.major_vessels_mask_ids = { - 6: "aorta", - 7: "inferior_vena_cava", - 17: "portal_vein_and_splenic_vein", - 58: "left_iliac_artery", - 59: "right_iliac_artery", - 60: "left_iliac_vena", - 61: "right_iliac_vena", - 110: "left_brachiocephalic_vena", - 111: "right_brachiocephalic_vena", - 112: "left_common_carotid_artery", - 113: "right_common_carotid_artery", - 119: "pulmonary_vein", - 123: "left_subclavian_artery", - 124: "right_subclavian_artery", - 125: "superior_vena_cava", - } - - self.lung_mask_ids = { - 28: "left_lung_upper_lobe", - 29: "left_lung_lower_lobe", - 30: "right_lung_upper_lobe", - 31: "right_lung_middle_lobe", - 32: "right_lung_lower_lobe", - 57: "trachea", - 132: "airway", - } - - self.bone_mask_ids = { - 33: "vertebrae_L5", - 34: "vertebrae_L4", - 35: "vertebrae_L3", - 36: "vertebrae_L2", - 37: "vertebrae_L1", - 38: "vertebrae_T12", - 39: "vertebrae_T11", - 40: "vertebrae_T10", - 41: "vertebrae_T9", - 42: "vertebrae_T8", - 43: "vertebrae_T7", - 44: "vertebrae_T6", - 45: "vertebrae_T5", - 46: "vertebrae_T4", - 47: "vertebrae_T3", - 48: "vertebrae_T2", - 49: "vertebrae_T1", - 50: "vertebrae_C7", - 51: "vertebrae_C6", - 52: "vertebrae_C5", - 53: "vertebrae_C4", - 54: "vertebrae_C3", - 55: "vertebrae_C2", - 56: "vertebrae_C1", - 63: "left_rib_1", - 64: "left_rib_2", - 65: "left_rib_3", - 66: "left_rib_4", - 67: "left_rib_5", - 68: "left_rib_6", - 69: "left_rib_7", - 70: "left_rib_8", - 71: "left_rib_9", - 72: "left_rib_10", - 73: "left_rib_11", - 74: "left_rib_12", - 75: "right_rib_1", - 76: "right_rib_2", - 77: "right_rib_3", - 78: "right_rib_4", - 79: "right_rib_5", - 80: "right_rib_6", - 81: "right_rib_7", - 82: "right_rib_8", - 83: "right_rib_9", - 84: "right_rib_10", - 85: "right_rib_11", - 86: "right_rib_12", - 87: "left_humerus", - 88: "right_humerus", - 89: "left_scapula", - 90: "right_scapula", - 91: "left_clavicula", - 92: "right_clavicula", - 93: "left_femur", - 94: "right_femur", - 95: "left_hip", - 96: "right_hip", - 114: "costal_cartilages", - 120: "skull", - 122: "sternum", - 127: "vertebrae_S1", - } - - self.soft_tissue_mask_ids = { - 121: "spinal_cord", - 118: "prostate", - 126: "thyroid_gland", - 62: "colon", - 19: "small_bowel", - 22: "brain", - 14: "left_kidney", - 15: "bladder", - 12: "stomach", - 13: "duodenum", - 8: "right_adrenal_gland", - 9: "left_adrenal_gland", - 10: "gallbladder", - 1: "liver", - 3: "spleen", - 4: "pancreas", - 5: "right_kidney", - 133: "soft_tissue", - } - - # From Base Class - # self.contrast_mask_ids = [135] - - self.set_other_and_all_mask_ids() - - def set_label_prompt(self, label_prompt: list[int]) -> None: - """ - Set specific anatomical structure labels to segment. - - Configures the segmentation to target specific anatomical structures - by their VISTA-3D label IDs instead of performing automatic segmentation. - - Args: - label_prompt (list): List of VISTA-3D anatomical structure IDs - to segment. See class attributes for available IDs - - Example: - >>> # Segment heart, aorta, and lungs only - >>> segmenter.set_label_prompt([115, 6, 28, 29, 30, 31, 32]) - >>> - >>> # Segment all cardiac structures - >>> heart_ids = list(segmenter.heart_mask_ids.keys()) - >>> segmenter.set_label_prompt(heart_ids) - """ - self.label_prompt = label_prompt - - def set_whole_image_segmentation(self) -> None: - """ - Configure for automatic whole-image segmentation. - - Resets the segmentation mode to automatic 'everything' segmentation, - clearing any previously set label prompts. - This is the default mode that segments all detectable structures. - - Example: - >>> # Reset to automatic segmentation after using label prompts - >>> segmenter.set_label_prompt([115, 6]) - >>> # ... perform label-prompted segmentation - >>> segmenter.set_whole_image_segmentation() # Reset to automatic - """ - self.label_prompt = None - - def segment_soft_tissue( - self, preprocessed_image: itk.image, labelmap_image: itk.image - ) -> itk.image: - """ - Add soft tissue segmentation to fill gaps in VISTA-3D output. - - VISTA-3D may not segment all tissue regions, leaving gaps between - structures. This method identifies soft tissue regions based on - intensity thresholds and adds them to the labelmap. - - Args: - preprocessed_image (itk.image): The preprocessed CT image - labelmap_image (itk.image): Existing VISTA-3D segmentation - - Returns: - itk.image: Updated labelmap with soft tissue regions filled - - Example: - >>> filled_labelmap = segmenter.segment_soft_tissue(preprocessed_image, vista_labelmap) - """ - hole_ids = [0] - labelmap_plus_soft_tissue_image = self.segment_connected_component( - preprocessed_image, - labelmap_image, - lower_threshold=-150, - upper_threshold=700, - labelmap_ids=hole_ids, - mask_id=list(self.soft_tissue_mask_ids.keys())[-1], - use_mid_slice=True, - ) - - return labelmap_plus_soft_tissue_image - - def preprocess_input(self, input_image: itk.image) -> itk.image: - """ - Preprocess the input image for VISTA-3D segmentation. - - Extends the base preprocessing with VISTA-3D specific adaptations: - - Adaptive spacing calculation based on input image properties - - No intensity rescaling (VISTA-3D handles raw CT intensities) - - Args: - input_image (itk.image): The input 3D CT image - - Returns: - itk.image: Preprocessed image optimized for VISTA-3D - - Note: - VISTA-3D works best with the original image spacing and intensity - values, so minimal preprocessing is applied. - """ - if self.target_spacing == 0.0: - spacing = input_image.GetSpacing() - self.target_spacing = (spacing[0] + spacing[1] + spacing[2]) / 3 - self.target_spacing = max(self.target_spacing, 0.5) - - preprocessed_image = super().preprocess_input(input_image) - - return preprocessed_image - - def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: - """ - Run VISTA-3D segmentation on the preprocessed image. - - Performs segmentation using the VISTA-3D model with the configured - interaction mode (automatic, point prompts, or label prompts). The - method handles model loading, inference, and post-processing including - soft tissue gap filling. - - Args: - preprocessed_image (itk.image): The preprocessed CT image ready - for VISTA-3D inference - - Returns: - itk.image: The segmentation labelmap with VISTA-3D labels and - additional soft tissue segmentation - - Raises: - ValueError: If no segmentation output is produced - RuntimeError: If model inference fails - - Note: - The method automatically selects the interaction mode based on - the configured prompts: - - No prompts: Everything segmentation - - Label prompt set: Specific structure segmentation - - Example: - >>> labelmap = segmenter.segmentation_method(preprocessed_ct) - """ - # Add bundle_path to sys.path only if not already present to avoid duplicates - if self.bundle_path not in sys.path: - sys.path.append(self.bundle_path) - - from hugging_face_pipeline import HuggingFacePipelineHelper - - hf_pipeline_helper = HuggingFacePipelineHelper(self.model_name) - hf_pipeline = hf_pipeline_helper.init_pipeline( - os.path.join(self.bundle_path, "vista3d_pretrained_model"), - device=self.device, - resample_spacing=( - self.target_spacing, - self.target_spacing, - self.target_spacing, - ), - ) - - # Use TemporaryDirectory context manager for exception-safe cleanup - with tempfile.TemporaryDirectory() as tmp_dir: - tmp_input_file_name = os.path.join(tmp_dir, "tmp.nii.gz") - itk.imwrite(preprocessed_image, tmp_input_file_name, compression=True) - - hf_inputs: list[dict[str, str | int | list[int]]] = [ - {"image": tmp_input_file_name} - ] - if self.label_prompt is None: - hf_inputs[0].update( - { - "label_prompt": hf_pipeline.EVERYTHING_LABEL, - } - ) - else: - hf_inputs[0].update( - { - "label_prompt": self.label_prompt, - } - ) - - hf_pipeline(hf_inputs, output_dir=tmp_dir) - - output_itk: Optional[itk.image] = None - for file_name in os.listdir(os.path.join(tmp_dir, "tmp")): - if file_name.endswith(".nii.gz"): - output_itk = itk.imread(os.path.join(tmp_dir, "tmp", file_name)) - output_itk.CopyInformation(preprocessed_image) - break - - if output_itk is None: - raise ValueError("No output image found") - - output_itk = self.segment_soft_tissue(preprocessed_image, output_itk) - - return output_itk diff --git a/src/physiomotion4d/segment_chest_vista_3d_nim.py b/src/physiomotion4d/segment_chest_vista_3d_nim.py deleted file mode 100644 index c5f76f2..0000000 --- a/src/physiomotion4d/segment_chest_vista_3d_nim.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Module for segmenting chest CT images using VISTA3D.""" - -# Please start vista3d docker: -# docker run --rm -it --name vista3d --runtime=nvidia -# -e CUDA_VISIBLE_DEVICES=0 -# -e NGC_API_KEY=$NGC_API_KEY -# --shm-size=8G -p 8000:8000 -# -v /tmp/data:/home/aylward/tmp/data nvcr.io/nim/nvidia/vista3d:latest - -import io -import json -import logging -import os -import socket -import tempfile -import zipfile -from urllib.error import HTTPError, URLError -from urllib.request import Request, urlopen - -import itk - -from physiomotion4d.segment_chest_vista_3d import SegmentChestVista3D - - -class SegmentChestVista3DNIM(SegmentChestVista3D): - """ - A class that inherits from physioSegmentChest and implements the - segmentation method using VISTA3D. - """ - - def __init__(self, log_level: int | str = logging.INFO) -> None: - """Initialize the vista3d class. - - Args: - log_level: Logging level (default: logging.INFO) - """ - super().__init__(log_level=log_level) - - self.invoke_url = "http://localhost:8000/v1/vista3d/inference" - self.wsl_docker_tmp_file = ( - "//wsl.localhost/Ubuntu/home/saylward/tmp/data/tmp.nii.gz" - ) - self.docker_tmp_file = "/tmp/data/tmp.nii.gz" - - def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: - """ - Run VISTA3D on the preprocessed image using the NIM and return result. - - Args: - preprocessed_image (itk.image): The preprocessed image to segment. - - Returns: - the segmented image. - """ - - # Post the image to file.io and get the link - itk.imwrite(preprocessed_image, self.wsl_docker_tmp_file, compression=True) - - payload = {"image": self.docker_tmp_file, "prompts": {}} - - # Call the API (stdlib HTTP client; avoids needing requests stubs) - payload_bytes = json.dumps(payload).encode("utf-8") - req = Request( - self.invoke_url, - data=payload_bytes, - headers={"Content-Type": "application/json"}, - method="POST", - ) - try: - # Use timeout to prevent indefinite hanging (300s = 5 minutes) - with urlopen(req, timeout=300) as resp: - response_content = resp.read() - except socket.timeout as e: - raise RuntimeError("VISTA3D NIM request timed out after 300 seconds") from e - except HTTPError as e: - raise RuntimeError( - f"VISTA3D NIM request failed: HTTP {e.code} {e.reason}" - ) from e - except URLError as e: - raise RuntimeError(f"VISTA3D NIM request failed: {e.reason}") from e - - # Get the result - labelmap_image = None - with tempfile.TemporaryDirectory() as temp_dir: - z = zipfile.ZipFile(io.BytesIO(response_content)) - z.extractall(temp_dir) - file_list = os.listdir(temp_dir) - for filename in file_list: - self.log_debug("Found file: %s", filename) - filepath = os.path.join(temp_dir, filename) - if os.path.isfile(filepath) and filename.endswith(".nii.gz"): - # SUCCESS: Return the results - labelmap_image = itk.imread(filepath, pixel_type=itk.SS) - break - - if labelmap_image is None: - raise Exception("Failed to get labelmap image from VISTA3D") - - # HERE - itk.imwrite(labelmap_image, "vista3d_labelmap.nii.gz", compression=True) - - # Include Soft Tissue - labelmap_image = self.segment_soft_tissue(preprocessed_image, labelmap_image) - - return labelmap_image diff --git a/src/physiomotion4d/workflow_convert_ct_to_vtk.py b/src/physiomotion4d/workflow_convert_ct_to_vtk.py index a522bfc..0d207a6 100644 --- a/src/physiomotion4d/workflow_convert_ct_to_vtk.py +++ b/src/physiomotion4d/workflow_convert_ct_to_vtk.py @@ -51,7 +51,6 @@ #: Supported segmentation backend identifiers. SEGMENTATION_METHODS: tuple[str, ...] = ( "total_segmentator", - "vista_3d", "simpleware_heart", ) @@ -63,7 +62,6 @@ class WorkflowConvertCTToVTK(PhysioMotion4DBase): - ``'total_segmentator'`` — :class:`SegmentChestTotalSegmentator` (CPU-capable, default). - - ``'vista_3d'`` — :class:`SegmentChestVista3D` (GPU-accelerated MONAI VISTA-3D). - ``'simpleware_heart'`` — :class:`SegmentHeartSimpleware` (cardiac only; requires a Simpleware Medical installation). @@ -106,8 +104,7 @@ def __init__( Args: segmentation_method: Segmentation backend to use. One of - ``'total_segmentator'`` (default), ``'vista_3d'``, or - ``'simpleware_heart'``. + ``'total_segmentator'`` (default) or ``'simpleware_heart'``. log_level: Logging level. Default: ``logging.INFO``. Raises: @@ -147,10 +144,6 @@ def _create_segmenter(self) -> SegmentAnatomyBase: ) return SegmentChestTotalSegmentator(log_level=self.log_level) - if self.segmentation_method_name == "vista_3d": - from physiomotion4d.segment_chest_vista_3d import SegmentChestVista3D - - return SegmentChestVista3D(log_level=self.log_level) if self.segmentation_method_name == "simpleware_heart": from physiomotion4d.segment_heart_simpleware import SegmentHeartSimpleware diff --git a/tests/GITHUB_WORKFLOWS.md b/tests/GITHUB_WORKFLOWS.md index 7e75674..25d42d0 100644 --- a/tests/GITHUB_WORKFLOWS.md +++ b/tests/GITHUB_WORKFLOWS.md @@ -218,12 +218,6 @@ These tests are **NOT** run in CI (even on GPU runners) because they require ext - Run locally: `pytest tests/test_segment_chest_total_segmentator.py -v -s` - Why excluded: Requires GPU, model inference -5. **VISTA-3D Tests** (`test_segment_chest_vista_3d.py`) - - Requires: CUDA GPU, VISTA-3D model weights - - Markers: `@pytest.mark.requires_data`, `@pytest.mark.slow` - - Run locally: `pytest tests/test_segment_chest_vista_3d.py -v -s` - - Why excluded: Requires GPU, model inference - **Why excluded**: - These tests take 5-15 minutes each, even with GPU acceleration - Registration algorithms are computationally intensive @@ -240,7 +234,7 @@ pytest tests/ -v -m "slow" pytest tests/test_register_images_ants.py tests/test_register_images_icon.py -v -s # Run only segmentation tests -pytest tests/test_segment_chest_total_segmentator.py tests/test_segment_chest_vista_3d.py -v -s +pytest tests/test_segment_chest_total_segmentator.py -v -s ``` **Scheduled slow tests**: diff --git a/tests/README.md b/tests/README.md index f9ca10c..50e95cd 100644 --- a/tests/README.md +++ b/tests/README.md @@ -19,7 +19,6 @@ This directory contains comprehensive test suites for the PhysioMotion4D package ### Segmentation Tests (GPU Required) - **`test_segment_chest_total_segmentator.py`** - TotalSegmentator chest CT segmentation -- **`test_segment_chest_vista_3d.py`** - NVIDIA VISTA-3D segmentation (requires 20GB+ RAM) ### Registration Tests (Slow ~5-10 min) - **`test_register_images_ants.py`** - ANTs deformable registration @@ -103,8 +102,7 @@ pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v -s -- ### Common Test Commands ```bash # Skip GPU-dependent tests -pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py \ - --ignore=tests/test_segment_chest_vista_3d.py +pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py # Run with coverage pytest tests/ --cov=src/physiomotion4d --cov-report=html @@ -146,8 +144,8 @@ test_convert_nrrd_4d_to_3d ↓ ├─→ test_register_images_icon ↓ ↓ test_segment_chest_total_segmentator ────→ test_contour_tools - ↓ ↓ -test_segment_chest_vista_3d test_convert_vtk_to_usd_polymesh + ↓ + test_convert_vtk_to_usd_polymesh ``` Fixtures in `conftest.py` automatically manage these dependencies. @@ -179,8 +177,8 @@ Tests automatically run on pull requests via GitHub Actions. The CI workflow: - Place `TruncalValve_4DCT.seq.nrrd` there to avoid download **Problem: Out of memory errors** -- VISTA-3D requires 20GB+ RAM -- Skip with: `pytest tests/ --ignore=tests/test_segment_chest_vista_3d.py` +- Some segmentation models require significant RAM +- Skip segmentation tests with: `pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py` **Problem: Test timeout** - Global timeout: 900 seconds (15 minutes) diff --git a/tests/TESTING_GUIDE.md b/tests/TESTING_GUIDE.md index 2c15246..ee316ef 100644 --- a/tests/TESTING_GUIDE.md +++ b/tests/TESTING_GUIDE.md @@ -6,7 +6,7 @@ This guide explains how to set up and use the PhysioMotion4D test suite. The test suite validates the complete PhysioMotion4D pipeline: - **Data download and conversion** - 4D NRRD to 3D time series -- **Segmentation** - TotalSegmentator and VISTA-3D chest CT segmentation +- **Segmentation** - TotalSegmentator chest CT segmentation - **Registration** - ANTs and ICON deformable registration - **Contour extraction** - PyVista mesh generation from segmentation masks - **USD conversion** - VTK to USD format for Omniverse @@ -80,8 +80,8 @@ test_convert_nrrd_4d_to_3d ↓ ├─→ test_register_images_icon ↓ ↓ test_segment_chest_total_segmentator ────→ test_contour_tools - ↓ ↓ -test_segment_chest_vista_3d test_convert_vtk_to_usd_polymesh + ↓ + test_convert_vtk_to_usd_polymesh ``` ### Test Markers @@ -97,7 +97,6 @@ Some tests are skipped or marked slow due to: **GPU-Dependent Tests** (skipped in CI): - `test_segment_chest_total_segmentator.py` - Requires GPU for inference -- `test_segment_chest_vista_3d.py` - Requires GPU + 20GB+ RAM **Computationally Intensive Tests**: - Registration tests (ANTs, ICON) - Marked slow, run locally only @@ -142,7 +141,6 @@ tests/ │ └── slice_001.mha └── results/ # Test outputs ├── segmentation_total_segmentator/ - ├── segmentation_vista3d/ ├── contour_tools/ ├── usd_polymesh/ ├── registration_ants/ @@ -164,15 +162,15 @@ tests/ **Manual fix**: Place `TruncalValve_4DCT.seq.nrrd` in `data/Slicer-Heart-CT/` -### Memory Errors (VISTA-3D Tests) +### Memory Errors (Segmentation Tests) -**Problem**: `RuntimeError: not enough memory: you tried to allocate 20GB` +**Problem**: `RuntimeError: not enough memory` -**Root Cause**: VISTA-3D requires full-resolution CT images, needs 20GB+ RAM +**Root Cause**: Segmentation models require significant RAM and GPU memory **Solutions**: -- Skip VISTA-3D tests: `pytest tests/ --ignore=tests/test_segment_chest_vista_3d.py` -- Run on system with 24GB+ RAM +- Skip segmentation tests: `pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py` +- Run on system with sufficient RAM and GPU memory - Tests are automatically skipped in CI ### Test Timeout @@ -205,7 +203,6 @@ size = (int(size_itk[0]), int(size_itk[1]), int(size_itk[2])) **Solution**: ✅ **Fixed!** Tests now use correct fixture names: - `segmenter_total_segmentator` for TotalSegmentator -- `segmenter_vista_3d` for VISTA-3D - `registrar_ants` for ANTs - `registrar_icon` for ICON @@ -329,9 +326,9 @@ pytest tests/test_usd_merge.py::TestUSDMerge::test_specific_test -v **Q: What if I don't have a GPU?** - Most tests run on CPU (slower but functional) -- Segmentation tests (TotalSegmentator, VISTA-3D) require GPU +- Segmentation tests (TotalSegmentator) require GPU - Registration tests (ANTs, ICON) benefit from GPU but work on CPU -- Skip GPU tests: `pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py --ignore=tests/test_segment_chest_vista_3d.py` +- Skip GPU tests: `pytest tests/ --ignore=tests/test_segment_chest_total_segmentator.py` **Q: How do I run tests without downloading data?** Place `TruncalValve_4DCT.seq.nrrd` in `data/Slicer-Heart-CT/` or `tests/data/Slicer-Heart-CT/` before running tests. The test will automatically detect and use it. @@ -341,7 +338,7 @@ Yes! Modify the `download_truncal_valve_data` fixture in `tests/conftest.py` to **Q: Why do some tests take so long?** - Registration (ANTs/ICON): 5-10 minutes each (deformable registration is computationally intensive) -- Segmentation (TotalSegmentator/VISTA-3D): 10-15 minutes (deep learning inference on full CT volumes) +- Segmentation (TotalSegmentator): 10-15 minutes (deep learning inference on full CT volumes) - Data download: First time only (~1.2GB file) - Everything else: <1 minute diff --git a/tests/conftest.py b/tests/conftest.py index 4bdda9f..24b142a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,7 @@ in the tests directory via pytest's automatic fixture discovery. """ +import os import shutil import urllib.request from datetime import datetime, timedelta @@ -21,7 +22,6 @@ 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_vista_3d import SegmentChestVista3D from physiomotion4d.segment_heart_simpleware import SegmentHeartSimpleware from physiomotion4d.transform_tools import TransformTools @@ -287,16 +287,21 @@ def download_truncal_valve_data(test_directories): return input_image_filename # Try to download if not found locally - input_image_url = "https://github.com/Slicer-Heart-CT/Slicer-Heart-CT/releases/download/TestingData/TruncalValve_4DCT.seq.nrrd" + input_image_url = "https://github.com/SlicerHeart/SlicerHeart/releases/download/TestingData/TruncalValve_4DCT.seq.nrrd" print(f"\nDownloading TruncalValve 4D CT data from {input_image_url}...") try: urllib.request.urlretrieve(input_image_url, str(input_image_filename)) print(f"Downloaded to {input_image_filename}") - except urllib.error.HTTPError as e: - pytest.skip( - f"Could not download test data: {e}. Please manually place TruncalValve_4DCT.seq.nrrd in {data_dir}" + except urllib.error.URLError as e: + msg = ( + f"Could not download test data: {e}. " + f"Please manually place TruncalValve_4DCT.seq.nrrd in {data_dir}" ) + if os.environ.get("CI"): + pytest.fail(msg) + else: + pytest.skip(msg) return input_image_filename @@ -381,12 +386,6 @@ def segmenter_total_segmentator(): return SegmentChestTotalSegmentator() -@pytest.fixture(scope="session") -def segmenter_vista_3d(): - """Create a SegmentChestVista3D instance.""" - return SegmentChestVista3D() - - @pytest.fixture(scope="session") def segmenter_simpleware(): """Create a SegmentHeartSimpleware instance.""" diff --git a/tests/test_experiments.py b/tests/test_experiments.py index 341f324..b858fe0 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -380,8 +380,6 @@ def test_experiment_heart_gated_ct_to_usd(): 3. 2-generate_segmentation.py (segments registered images) 4. 3-transform_dynamic_and_static_contours.py (transforms segmentations) 5. 4-merge_dynamic_and_static_usd.py (merges into final USD) - 6. test_vista3d_class.py (tests segmentation class) - 7. test_vista3d_inMem.py (tests in-memory segmentation) Each script depends on outputs from previous scripts. Execution stops on first failure to prevent cascading errors. diff --git a/tests/test_segment_chest_vista_3d.py b/tests/test_segment_chest_vista_3d.py deleted file mode 100644 index dc88884..0000000 --- a/tests/test_segment_chest_vista_3d.py +++ /dev/null @@ -1,312 +0,0 @@ -#!/usr/bin/env python -""" -Test for chest CT segmentation using VISTA-3D. - -This test depends on test_convert_nrrd_4d_to_3d and tests segmentation -functionality on two time points from the converted 3D data. -""" - -import itk -import numpy as np -import pytest - - -@pytest.mark.requires_data -@pytest.mark.slow -class TestSegmentChestVista3D: - """Test suite for VISTA-3D chest CT segmentation.""" - - def test_segmenter_initialization(self, segmenter_vista_3d): - """Test that SegmentChestVista3D initializes correctly.""" - assert segmenter_vista_3d is not None, "Segmenter not initialized" - assert segmenter_vista_3d.device is not None, "CUDA device not initialized" - - # Check that anatomical structure ID mappings are defined - assert len(segmenter_vista_3d.heart_mask_ids) > 0, "Heart mask IDs not defined" - assert len(segmenter_vista_3d.major_vessels_mask_ids) > 0, ( - "Major vessels mask IDs not defined" - ) - assert len(segmenter_vista_3d.lung_mask_ids) > 0, "Lung mask IDs not defined" - assert len(segmenter_vista_3d.bone_mask_ids) > 0, "Bone mask IDs not defined" - assert len(segmenter_vista_3d.soft_tissue_mask_ids) > 0, ( - "Soft tissue mask IDs not defined" - ) - - # Check VISTA-3D specific attributes - assert segmenter_vista_3d.bundle_path is not None, "Bundle path not set" - assert segmenter_vista_3d.label_prompt is None, ( - "Label prompt should be None initially" - ) - - print("\n✓ Segmenter initialized with correct parameters") - print(f" Heart structures: {len(segmenter_vista_3d.heart_mask_ids)}") - print(f" Major vessels: {len(segmenter_vista_3d.major_vessels_mask_ids)}") - print(f" Lung structures: {len(segmenter_vista_3d.lung_mask_ids)}") - print(f" Bone structures: {len(segmenter_vista_3d.bone_mask_ids)}") - print( - f" Soft tissue structures: {len(segmenter_vista_3d.soft_tissue_mask_ids)}" - ) - print(f" Bundle path: {segmenter_vista_3d.bundle_path}") - - def test_segment_single_image( - self, segmenter_vista_3d, test_images, test_directories - ): - """Test automatic segmentation on a single time point.""" - output_dir = test_directories["output"] - - # Ensure we're in automatic segmentation mode - segmenter_vista_3d.set_whole_image_segmentation() - - # Test on first time point only - input_image = test_images[0] - - print("\nSegmenting time point 0 (automatic mode)...") - print(f" Input image size: {itk.size(input_image)}") - - # Run segmentation - result = segmenter_vista_3d.segment(input_image, contrast_enhanced_study=False) - - # Verify result is a dictionary with expected keys - assert isinstance(result, dict), "Result should be a dictionary" - expected_keys = [ - "labelmap", - "lung", - "heart", - "major_vessels", - "bone", - "soft_tissue", - "other", - "contrast", - ] - for key in expected_keys: - assert key in result, f"Missing key '{key}' in result" - assert result[key] is not None, f"Result['{key}'] is None" - - # Verify labelmap properties - labelmap = result["labelmap"] - assert itk.size(labelmap) == itk.size(input_image), "Labelmap size mismatch" - - # Check that labels are present - labelmap_arr = itk.array_from_image(labelmap) - unique_labels = np.unique(labelmap_arr) - assert len(unique_labels) > 1, "Labelmap should contain multiple labels" - - print("✓ Segmentation complete for time point 0") - print(f" Labelmap size: {itk.size(labelmap)}") - print(f" Unique labels: {len(unique_labels)}") - - # Save results - seg_output_dir = output_dir / "segmentation_vista3d" - seg_output_dir.mkdir(exist_ok=True) - - itk.imwrite( - labelmap, str(seg_output_dir / "slice_000_labelmap.mha"), compression=True - ) - print(f" Saved labelmap to: {seg_output_dir / 'slice_000_labelmap.mha'}") - - def test_segment_multiple_images( - self, segmenter_vista_3d, test_images, test_directories - ): - """Test automatic segmentation on two time points.""" - output_dir = test_directories["output"] - seg_output_dir = output_dir / "segmentation_vista3d" - seg_output_dir.mkdir(exist_ok=True) - - # Ensure automatic segmentation mode - segmenter_vista_3d.set_whole_image_segmentation() - - results = [] - for i, input_image in enumerate(test_images[0:2]): - print(f"\nSegmenting time point {i}...") - - result = segmenter_vista_3d.segment( - input_image, contrast_enhanced_study=False - ) - results.append(result) - - # Save labelmap for each time point - labelmap = result["labelmap"] - output_file = seg_output_dir / f"slice_{i:03d}_labelmap.mha" - itk.imwrite(labelmap, str(output_file), compression=True) - - print(f"✓ Time point {i} complete") - print(f" Saved to: {output_file}") - - assert len(results) == 2, "Expected 2 segmentation results" - print(f"\n✓ Successfully segmented {len(results)} time points") - - def test_anatomy_group_masks(self, segmenter_vista_3d, test_images): - """Test that anatomy group masks are created correctly.""" - segmenter_vista_3d.set_whole_image_segmentation() - input_image = test_images[0] - - # Run segmentation - result = segmenter_vista_3d.segment(input_image, contrast_enhanced_study=False) - - # Check each anatomy group mask - anatomy_groups = [ - "lung", - "heart", - "major_vessels", - "bone", - "soft_tissue", - "other", - ] - - for group in anatomy_groups: - mask = result[group] - assert mask is not None, f"{group} mask is None" - - # Check that mask is binary - mask_arr = itk.array_from_image(mask) - unique_values = np.unique(mask_arr) - assert len(unique_values) <= 2, f"{group} mask should be binary" - assert 0 in unique_values, f"{group} mask should contain background" - - # Check that mask has same size as input - assert itk.size(mask) == itk.size(input_image), ( - f"{group} mask size mismatch" - ) - - print("\n✓ All anatomy group masks created correctly") - for group in anatomy_groups: - mask_arr = itk.array_from_image(result[group]) - num_voxels = np.sum(mask_arr > 0) - print(f" {group}: {num_voxels} voxels") - - def test_label_prompt_segmentation( - self, segmenter_vista_3d, test_images, test_directories - ): - """Test segmentation with specific label prompts.""" - output_dir = test_directories["output"] - seg_output_dir = output_dir / "segmentation_vista3d" - seg_output_dir.mkdir(exist_ok=True) - - input_image = test_images[0] - - # Test with heart and aorta labels only - heart_aorta_labels = [115, 6] # Heart and aorta - segmenter_vista_3d.set_label_prompt(heart_aorta_labels) - - print(f"\nSegmenting with label prompts: {heart_aorta_labels}") - result = segmenter_vista_3d.segment(input_image, contrast_enhanced_study=False) - - # Verify result - assert isinstance(result, dict), "Result should be a dictionary" - labelmap = result["labelmap"] - - # Check that only prompted labels are present (plus background and soft tissue fill) - labelmap_arr = itk.array_from_image(labelmap) - unique_labels = np.unique(labelmap_arr) - - print("✓ Label prompt segmentation complete") - print(f" Unique labels: {unique_labels}") - - # Save result - output_file = seg_output_dir / "slice_000_label_prompt.mha" - itk.imwrite(labelmap, str(output_file), compression=True) - print(f" Saved to: {output_file}") - - # Reset to whole image segmentation - segmenter_vista_3d.set_whole_image_segmentation() - - def test_contrast_detection(self, segmenter_vista_3d, test_images): - """Test contrast detection functionality.""" - segmenter_vista_3d.set_whole_image_segmentation() - input_image = test_images[0] - - # Test without contrast - result_no_contrast = segmenter_vista_3d.segment( - input_image, contrast_enhanced_study=False - ) - contrast_mask_no = result_no_contrast["contrast"] - - # Test with contrast flag - result_with_contrast = segmenter_vista_3d.segment( - input_image, contrast_enhanced_study=True - ) - contrast_mask_yes = result_with_contrast["contrast"] - - # 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("\n✓ Contrast 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") - - def test_preprocessing(self, segmenter_vista_3d, test_images): - """Test preprocessing functionality.""" - segmenter_vista_3d.set_whole_image_segmentation() - input_image = test_images[0] - - # Get original properties - original_spacing = itk.spacing(input_image) - - # Preprocessing is done internally by segment(), not exposed as public method - # Just verify that segment() works (which includes preprocessing) - result = segmenter_vista_3d.segment(input_image, contrast_enhanced_study=False) - - # Check that segmentation was successful (which means preprocessing worked) - assert result is not None, "Segmentation result is None" - assert "labelmap" in result, "Labelmap not in result" - - print("\n✓ Preprocessing tested (via successful segmentation)") - print(f" Original image spacing: {original_spacing}") - - def test_postprocessing(self, segmenter_vista_3d, test_images): - """Test postprocessing functionality.""" - segmenter_vista_3d.set_whole_image_segmentation() - input_image = test_images[0] - - # Run full segmentation to get labelmap - result = segmenter_vista_3d.segment(input_image, contrast_enhanced_study=False) - labelmap = result["labelmap"] - - # Postprocessing is part of segment(), verify output is properly sized - assert itk.size(labelmap) == itk.size(input_image), ( - "Postprocessing failed: size mismatch" - ) - - # Check that labelmap has been resampled to original spacing - original_spacing = itk.spacing(input_image) - labelmap_spacing = itk.spacing(labelmap) - - # Spacing should match (within floating point tolerance) - for i in range(3): - assert abs(labelmap_spacing[i] - original_spacing[i]) < 0.01, ( - f"Spacing mismatch at dimension {i}" - ) - - print("\n✓ Postprocessing tested") - print(f" Original spacing: {original_spacing}") - print(f" Labelmap spacing: {labelmap_spacing}") - - def test_set_and_reset_prompts(self, segmenter_vista_3d): - """Test setting and resetting label prompt mode.""" - # Initially should be in automatic mode - assert segmenter_vista_3d.label_prompt is None, ( - "Label prompt should be None initially" - ) - - # Set label prompt - segmenter_vista_3d.set_label_prompt([115, 6]) - assert segmenter_vista_3d.label_prompt == [ - 115, - 6, - ], "Label prompt not set correctly" - - # Reset to whole image - segmenter_vista_3d.set_whole_image_segmentation() - assert segmenter_vista_3d.label_prompt is None, ( - "Label prompt should be None after reset" - ) - - print("\n✓ Prompt setting and resetting works correctly") - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"])