ENH: Update fit-to-patient for non-contrast heart models#87
Conversation
Non-contrast heart models from TotalSegmentator are incomplete, but this update makes registration to a statistical model possible.
|
Warning Review limit reached
Next review available in: 35 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughThe changes rename workflow execution methods to ChangesSegmentation pipeline
Workflow API migration
Workflow utility updates
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant WorkflowClient
participant WorkflowFitStatisticalModelToPatient
participant WorkflowReconstructHighres4DCT
participant ResultWriters
WorkflowClient->>WorkflowFitStatisticalModelToPatient: Configure PCA and call process()
WorkflowFitStatisticalModelToPatient-->>WorkflowClient: Return registration result
WorkflowClient->>WorkflowReconstructHighres4DCT: Call process()
WorkflowReconstructHighres4DCT-->>WorkflowClient: Return reconstructions and transforms
WorkflowClient->>ResultWriters: Save models, images, and transforms
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/physiotwin4d/workflow_convert_image_to_usd.py (1)
83-84: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUpdate the docstring to reflect the new default value.
The default value for
frames_per_secondwas changed to1.0(at line 57), but the docstring still states it defaults to24.0. As per coding guidelines, docstrings must be updated and kept factual for changed public methods.📝 Proposed fix
- frames_per_second: Frames per second for animated USD time series. - Defaults to 24.0, matching the underlying VTK-to-USD converter. + frames_per_second: Frames per second for animated USD time series. + Defaults to 1.0.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/physiotwin4d/workflow_convert_image_to_usd.py` around lines 83 - 84, Update the frames_per_second parameter documentation in the workflow conversion function’s docstring to state that the default is 1.0 instead of 24.0, matching the value defined in the function signature.Source: Coding guidelines
docs/developer/segmentation.rst (1)
40-44: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDoc references stale
contrastid 135.
segment_chest_total_segmentator_with_contrast.pynow registers"contrast"at id 155 (moved from 135 to avoid colliding with the newbody_extremities=135 label), but this line still documents "contrast at id 135".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/developer/segmentation.rst` around lines 40 - 44, Update the taxonomy documentation for the default “contrast” placeholder to use id 155 instead of 135, while retaining the “soft_tissue” id 133 reference and the surrounding description.src/physiotwin4d/segment_chest_total_segmentator.py (1)
46-47: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRestore
target_spacingto1.5or update the contract around resampling.self.target_spacing = 0.0disablesSegmentAnatomyBase.preprocess_input(), contradicts the class docstring, and breaks the current initialization test that expects1.5. If the new behavior is intentional, update the docstring and test to match.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/physiotwin4d/segment_chest_total_segmentator.py` around lines 46 - 47, Restore self.target_spacing to 1.5 in the TotalSegmentator class so it matches the documented contract and existing initialization test, preserving the expected resampling behavior through SegmentAnatomyBase.preprocess_input().
🧹 Nitpick comments (3)
src/physiotwin4d/contour_tools.py (1)
34-35: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate docstrings for the changed public methods.
The static-method conversion is compatible with the supplied callers, but the public documentation is incomplete:
transform_contoursomits several parameters,create_reference_imagelacks argument and return documentation, andcreate_mask_from_meshhas no docstring. Please update all five changed methods with complete, factual docstrings.As per coding guidelines, update docstrings for every changed public method.
Also applies to: 74-75, 138-139, 167-168, 363-364
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/physiotwin4d/contour_tools.py` around lines 34 - 35, Update the docstrings for all five changed public methods, including extract_contours, transform_contours, create_reference_image, create_mask_from_mesh, and the method at the additional referenced location. Document every parameter and the return value with accurate descriptions, ensuring transform_contours includes its omitted parameters and create_reference_image and create_mask_from_mesh have complete argument and return documentation.Source: Coding guidelines
src/physiotwin4d/segment_anatomy_base.py (1)
456-463: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winGroup labelmaps aren't cast to uint8 like the main
"labelmap"key.
labelmap_imageis explicitly re-cast touint8before being returned, but thelabelmapsdict (heart/lung/etc.) is built from the pre-cast image and returned as-is. All current segmenters happen to already produce uint8 labelmaps, so this is not currently observable, but the returned dict has no guaranteed-consistent pixel type across its values, which downstream consumers (writers, comparisons) may implicitly assume.♻️ Suggested fix
labelmaps = self.create_anatomy_group_labelmaps(labelmap_image) labelmap_image = itk.GetImageFromArray( itk.GetArrayFromImage(labelmap_image).astype(np.uint8) ) labelmap_image.CopyInformation(input_image) - return {"labelmap": labelmap_image, **labelmaps} + labelmaps = { + group: itk.GetImageFromArray( + itk.GetArrayFromImage(group_labelmap).astype(np.uint8) + ) + for group, group_labelmap in labelmaps.items() + } + for group_labelmap, source in zip(labelmaps.values(), labelmaps.values()): + group_labelmap.CopyInformation(input_image) + + return {"labelmap": labelmap_image, **labelmaps}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/physiotwin4d/segment_anatomy_base.py` around lines 456 - 463, Update the labelmap construction flow around create_anatomy_group_labelmaps so every group labelmap in labelmaps is converted to uint8 before returning. Apply the same pixel-type normalization used for the main "labelmap" while preserving each group's existing image metadata and the return structure.tutorials/tutorial_04_fit_statistical_model_to_patient.py (1)
101-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPersisted ITK images/transforms omit
compression=Truein two tutorials. Bothitk.imwriteanditk.transformwritedefault theircompressionparameter toFalse; the coding guidelines require compressed storage for both ITK images and.hdftransforms.
tutorials/tutorial_04_fit_statistical_model_to_patient.py#L101-L110: addcompression=Trueto the threeitk.imwrite(...)calls (patient image, patient labelmap, heart labelmap).tutorials/tutorial_06_reconstruct_highres_4d_ct.py#L107-L111: addcompression=Trueto the twoitk.transformwrite(...)calls (forward and inverse transform).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tutorials/tutorial_04_fit_statistical_model_to_patient.py` around lines 101 - 110, Enable compression on all persisted ITK outputs: update the three itk.imwrite calls in tutorials/tutorial_04_fit_statistical_model_to_patient.py for patient_image, patient_labelmap, and heart_labelmap, and the two itk.transformwrite calls in tutorials/tutorial_06_reconstruct_highres_4d_ct.py for the forward and inverse transforms, by passing compression=True.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@experiments/Heart-Create_Statistical_Model/README.md`:
- Line 164: Update the workflow result assignment around
WorkflowFitStatisticalModelToPatient.process() so it selects
registered_template_model_surface from the returned dictionary instead of
assigning the entire dictionary to registered_mesh.
In `@src/physiotwin4d/segment_chest_total_segmentator.py`:
- Line 286: Remove the three unconditional debug itk.imwrite calls for
preprocessed_image, labelmap_arr_heart, and labelmap_arr_body from the
segmentation flow. Do not replace them with other production writes; preserve
the segmentation behavior while eliminating the working-directory artifacts and
associated debug image creation if it is used only by those writes.
- Around line 369-393: Binarize the heart interior mask before constructing
interior_image in the high-resolution heart expansion block: make interior_arr
contain 1 for labels 141–144 and 0 elsewhere, while preserving the existing
dilation and erosion calls using foreground value 1. Keep the subsequent
exterior mask application and labelmap replacement behavior unchanged.
In `@src/physiotwin4d/workflow_convert_image_to_usd.py`:
- Line 228: Wrap the assignment to self.reference_segmentation_results in the
workflow conversion code so each resulting line is no longer than 88 characters,
preserving the existing segmenter.segment(self.reference_image) behavior.
In `@tests/conftest.py`:
- Around line 517-529: Remove the unconditional "contrast" lookup from the
fixture result assembled with segmenter_total_segmentator, since
create_anatomy_group_labelmaps only returns registered groups. Keep the existing
groups unchanged; if contrast output is required, use the
segmenter_total_segmentator_with_contrast fixture instead.
In `@tutorials/tutorial_04_fit_statistical_model_to_patient.py`:
- Around line 128-143: Guard PCA registration in the workflow setup using the
loaded pca_model: call set_use_pca_registration only when pca_model is not None,
while preserving the existing enabled configuration when it is available. Update
the result-saving flow around workflow.pca_coefficients so it does not call
tolist() when PCA was skipped, and save registered coefficients only when they
exist.
---
Outside diff comments:
In `@docs/developer/segmentation.rst`:
- Around line 40-44: Update the taxonomy documentation for the default
“contrast” placeholder to use id 155 instead of 135, while retaining the
“soft_tissue” id 133 reference and the surrounding description.
In `@src/physiotwin4d/segment_chest_total_segmentator.py`:
- Around line 46-47: Restore self.target_spacing to 1.5 in the TotalSegmentator
class so it matches the documented contract and existing initialization test,
preserving the expected resampling behavior through
SegmentAnatomyBase.preprocess_input().
In `@src/physiotwin4d/workflow_convert_image_to_usd.py`:
- Around line 83-84: Update the frames_per_second parameter documentation in the
workflow conversion function’s docstring to state that the default is 1.0
instead of 24.0, matching the value defined in the function signature.
---
Nitpick comments:
In `@src/physiotwin4d/contour_tools.py`:
- Around line 34-35: Update the docstrings for all five changed public methods,
including extract_contours, transform_contours, create_reference_image,
create_mask_from_mesh, and the method at the additional referenced location.
Document every parameter and the return value with accurate descriptions,
ensuring transform_contours includes its omitted parameters and
create_reference_image and create_mask_from_mesh have complete argument and
return documentation.
In `@src/physiotwin4d/segment_anatomy_base.py`:
- Around line 456-463: Update the labelmap construction flow around
create_anatomy_group_labelmaps so every group labelmap in labelmaps is converted
to uint8 before returning. Apply the same pixel-type normalization used for the
main "labelmap" while preserving each group's existing image metadata and the
return structure.
In `@tutorials/tutorial_04_fit_statistical_model_to_patient.py`:
- Around line 101-110: Enable compression on all persisted ITK outputs: update
the three itk.imwrite calls in
tutorials/tutorial_04_fit_statistical_model_to_patient.py for patient_image,
patient_labelmap, and heart_labelmap, and the two itk.transformwrite calls in
tutorials/tutorial_06_reconstruct_highres_4d_ct.py for the forward and inverse
transforms, by passing compression=True.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: abdc6a98-e2d4-4726-9063-94a05139b70e
📒 Files selected for processing (23)
docs/api/model_registration/index.rstdocs/api/workflows.rstdocs/developer/registration_models.rstdocs/developer/segmentation.rstdocs/tutorials.rstexperiments/Heart-Create_Statistical_Model/README.mdexperiments/Heart-GatedCT_To_USD/2-generate_segmentation.pyexperiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.pysrc/physiotwin4d/cli/fit_statistical_model_to_patient.pysrc/physiotwin4d/cli/reconstruct_highres_4d_ct.pysrc/physiotwin4d/contour_tools.pysrc/physiotwin4d/segment_anatomy_base.pysrc/physiotwin4d/segment_chest_total_segmentator.pysrc/physiotwin4d/segment_chest_total_segmentator_with_contrast.pysrc/physiotwin4d/segment_heart_simpleware.pysrc/physiotwin4d/workflow_convert_image_to_usd.pysrc/physiotwin4d/workflow_fit_statistical_model_to_patient.pysrc/physiotwin4d/workflow_reconstruct_highres_4d_ct.pytests/conftest.pytests/test_segment_chest_total_segmentator.pytutorials/tutorial_04_fit_statistical_model_to_patient.pytutorials/tutorial_06_reconstruct_highres_4d_ct.pytutorials/tutorial_08_cardiac_fit_model.py
| ) | ||
|
|
||
| registered_mesh = workflow.run_workflow() | ||
| registered_mesh = workflow.process() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Select the mesh from the workflow result dictionary.
WorkflowFitStatisticalModelToPatient.process() returns a dictionary, not a mesh directly. Assigning it to registered_mesh makes this example incorrect for callers expecting a mesh.
Suggested fix
-registered_mesh = workflow.process()
+result = workflow.process()
+registered_mesh = result["registered_template_model"]Use registered_template_model_surface instead if this experiment needs the surface mesh.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@experiments/Heart-Create_Statistical_Model/README.md` at line 164, Update the
workflow result assignment around WorkflowFitStatisticalModelToPatient.process()
so it selects registered_template_model_surface from the returned dictionary
instead of assigning the entire dictionary to registered_mesh.
| # Add heart around interior regions. | ||
| if self.has_highres_heart_license: | ||
| interior_mask = np.isin(labelmap_arr, [141, 142, 143, 144]) | ||
| interior_arr = labelmap_arr * interior_mask | ||
| interior_image = itk.GetImageFromArray(interior_arr.astype(np.uint8)) | ||
| interior_image.CopyInformation(preprocessed_image) | ||
| imMath = ImageTools() | ||
| spacing = interior_image.GetSpacing() | ||
| exterior_image = imMath.binary_dilate_image( | ||
| interior_image, round(7 / spacing[0]), 1, 0 | ||
| ) | ||
| exterior_image = imMath.binary_erode_image( | ||
| exterior_image, round(4 / spacing[0]), 1, 0 | ||
| ) | ||
| exterior_arr = itk.GetArrayFromImage(exterior_image) | ||
| mask_id = 51 # Heart mask id | ||
| exterior_arr = exterior_arr * mask_id | ||
| labelmap_arr = np.where( | ||
| labelmap_arr == 0, exterior_arr, labelmap_arr | ||
| ) | ||
| replace_arr = np.where(labelmap_arr == 133, exterior_arr, 0) | ||
| labelmap_arr = np.where( | ||
| replace_arr > 0, exterior_arr, labelmap_arr | ||
| ) | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Highres-heart expansion is a no-op: interior_arr is never binarized before dilation with foreground_value=1.
interior_arr = labelmap_arr * interior_mask keeps original label values {0, 141, 142, 143, 144}, but imMath.binary_dilate_image(interior_image, ..., 1, 0) / binary_erode_image(..., 1, 0) treat pixel value 1 as foreground (per ImageTools.binary_dilate_image/binary_erode_image, which sets SetForegroundValue(foreground_value) on BinaryDilateImageFilter/BinaryErodeImageFilter). Since interior_image never contains the value 1, the filters find zero foreground voxels, exterior_arr ends up all zeros, and the entire "add heart around interior regions" block silently does nothing — this appears to be the exact mechanism this PR needs for incomplete non-contrast heart models.
Compare with SegmentHeartSimpleware.segmentation_method, which correctly builds a true binary interior_array (tmp_array = (mask_array > 128).astype(np.uint8)) before the equivalent dilate/erode calls.
🐛 Suggested fix
interior_mask = np.isin(labelmap_arr, [141, 142, 143, 144])
- interior_arr = labelmap_arr * interior_mask
+ interior_arr = interior_mask.astype(np.uint8)
interior_image = itk.GetImageFromArray(interior_arr.astype(np.uint8))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Add heart around interior regions. | |
| if self.has_highres_heart_license: | |
| interior_mask = np.isin(labelmap_arr, [141, 142, 143, 144]) | |
| interior_arr = labelmap_arr * interior_mask | |
| interior_image = itk.GetImageFromArray(interior_arr.astype(np.uint8)) | |
| interior_image.CopyInformation(preprocessed_image) | |
| imMath = ImageTools() | |
| spacing = interior_image.GetSpacing() | |
| exterior_image = imMath.binary_dilate_image( | |
| interior_image, round(7 / spacing[0]), 1, 0 | |
| ) | |
| exterior_image = imMath.binary_erode_image( | |
| exterior_image, round(4 / spacing[0]), 1, 0 | |
| ) | |
| exterior_arr = itk.GetArrayFromImage(exterior_image) | |
| mask_id = 51 # Heart mask id | |
| exterior_arr = exterior_arr * mask_id | |
| labelmap_arr = np.where( | |
| labelmap_arr == 0, exterior_arr, labelmap_arr | |
| ) | |
| replace_arr = np.where(labelmap_arr == 133, exterior_arr, 0) | |
| labelmap_arr = np.where( | |
| replace_arr > 0, exterior_arr, labelmap_arr | |
| ) | |
| # Add heart around interior regions. | |
| if self.has_highres_heart_license: | |
| interior_mask = np.isin(labelmap_arr, [141, 142, 143, 144]) | |
| interior_arr = interior_mask.astype(np.uint8) | |
| interior_image = itk.GetImageFromArray(interior_arr.astype(np.uint8)) | |
| interior_image.CopyInformation(preprocessed_image) | |
| imMath = ImageTools() | |
| spacing = interior_image.GetSpacing() | |
| exterior_image = imMath.binary_dilate_image( | |
| interior_image, round(7 / spacing[0]), 1, 0 | |
| ) | |
| exterior_image = imMath.binary_erode_image( | |
| exterior_image, round(4 / spacing[0]), 1, 0 | |
| ) | |
| exterior_arr = itk.GetArrayFromImage(exterior_image) | |
| mask_id = 51 # Heart mask id | |
| exterior_arr = exterior_arr * mask_id | |
| labelmap_arr = np.where( | |
| labelmap_arr == 0, exterior_arr, labelmap_arr | |
| ) | |
| replace_arr = np.where(labelmap_arr == 133, exterior_arr, 0) | |
| labelmap_arr = np.where( | |
| replace_arr > 0, exterior_arr, labelmap_arr | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/physiotwin4d/segment_chest_total_segmentator.py` around lines 369 - 393,
Binarize the heart interior mask before constructing interior_image in the
high-resolution heart expansion block: make interior_arr contain 1 for labels
141–144 and 0 elsewhere, while preserving the existing dilation and erosion
calls using foreground value 1. Keep the subsequent exterior mask application
and labelmap replacement behavior unchanged.
|
|
||
| self.reference_segmentation = self.segmenter.segment(self.reference_image) | ||
| labelmap = self.reference_segmentation["labelmap"] | ||
| self.reference_segmentation_results = self.segmenter.segment(self.reference_image) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Keep lines at or below 88 characters.
This line is 90 characters long, which exceeds the maximum allowed length of 88 characters. As per coding guidelines, lines must be kept at or below 88 characters.
📝 Proposed fix
- self.reference_segmentation_results = self.segmenter.segment(self.reference_image)
+ self.reference_segmentation_results = self.segmenter.segment(
+ self.reference_image
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.reference_segmentation_results = self.segmenter.segment(self.reference_image) | |
| self.reference_segmentation_results = self.segmenter.segment( | |
| self.reference_image | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/physiotwin4d/workflow_convert_image_to_usd.py` at line 228, Wrap the
assignment to self.reference_segmentation_results in the workflow conversion
code so each resulting line is no longer than 88 characters, preserving the
existing segmenter.segment(self.reference_image) behavior.
Source: Coding guidelines
| labelmap = itk.imread(str(labelmap_file)) | ||
| masks = segmenter_total_segmentator.create_anatomy_group_masks(labelmap) | ||
| labelmaps = segmenter_total_segmentator.create_anatomy_group_labelmaps(labelmap) | ||
| results.append( | ||
| { | ||
| "labelmap": labelmap, | ||
| "lung": masks["lung"], | ||
| "heart": masks["heart"], | ||
| "major_vessels": masks["major_vessels"], | ||
| "bone": masks["bone"], | ||
| "soft_tissue": masks["soft_tissue"], | ||
| "other": masks["other"], | ||
| "contrast": masks["contrast"], | ||
| "lung": labelmaps["lung"], | ||
| "heart": labelmaps["heart"], | ||
| "major_vessels": labelmaps["major_vessels"], | ||
| "bone": labelmaps["bone"], | ||
| "soft_tissue": labelmaps["soft_tissue"], | ||
| "other": labelmaps["other"], | ||
| "contrast": labelmaps["contrast"], | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
labelmaps["contrast"] will raise KeyError — segmenter_total_segmentator has no "contrast" group.
This fixture is parametrized with the plain segmenter_total_segmentator (SegmentChestTotalSegmentator), which never registers a "contrast" taxonomy group — only SegmentChestTotalSegmentatorWithContrast does via _add_extra_taxonomy_groups(). create_anatomy_group_labelmaps() only returns keys present in taxonomy.group_names() (plus "other"), so labelmaps["contrast"] will raise KeyError for every invocation of this fixture.
🐛 Suggested fix
- "contrast": labelmaps["contrast"],
+ "contrast": labelmaps.get("contrast"),Or, if contrast results are actually required here, switch this fixture to use segmenter_total_segmentator_with_contrast instead.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| labelmap = itk.imread(str(labelmap_file)) | |
| masks = segmenter_total_segmentator.create_anatomy_group_masks(labelmap) | |
| labelmaps = segmenter_total_segmentator.create_anatomy_group_labelmaps(labelmap) | |
| results.append( | |
| { | |
| "labelmap": labelmap, | |
| "lung": masks["lung"], | |
| "heart": masks["heart"], | |
| "major_vessels": masks["major_vessels"], | |
| "bone": masks["bone"], | |
| "soft_tissue": masks["soft_tissue"], | |
| "other": masks["other"], | |
| "contrast": masks["contrast"], | |
| "lung": labelmaps["lung"], | |
| "heart": labelmaps["heart"], | |
| "major_vessels": labelmaps["major_vessels"], | |
| "bone": labelmaps["bone"], | |
| "soft_tissue": labelmaps["soft_tissue"], | |
| "other": labelmaps["other"], | |
| "contrast": labelmaps["contrast"], | |
| } | |
| labelmap = itk.imread(str(labelmap_file)) | |
| labelmaps = segmenter_total_segmentator.create_anatomy_group_labelmaps(labelmap) | |
| results.append( | |
| { | |
| "labelmap": labelmap, | |
| "lung": labelmaps["lung"], | |
| "heart": labelmaps["heart"], | |
| "major_vessels": labelmaps["major_vessels"], | |
| "bone": labelmaps["bone"], | |
| "soft_tissue": labelmaps["soft_tissue"], | |
| "other": labelmaps["other"], | |
| "contrast": labelmaps.get("contrast"), | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/conftest.py` around lines 517 - 529, Remove the unconditional
"contrast" lookup from the fixture result assembled with
segmenter_total_segmentator, since create_anatomy_group_labelmaps only returns
registered groups. Keep the existing groups unchanged; if contrast output is
required, use the segmenter_total_segmentator_with_contrast fixture instead.
| workflow.set_use_pca_registration( | ||
| use_pca_registration=True, | ||
| pca_model=pca_model, | ||
| use_surface=False, | ||
| ) | ||
| if pca_model is not None: | ||
| workflow.set_use_pca_registration(True, pca_model=pca_model) | ||
|
|
||
| # %% | ||
| # Workflow execution | ||
| result = workflow.run_workflow() | ||
| workflow_results = workflow.process() | ||
|
|
||
| # %% | ||
| # Result saving | ||
| registered_surface: pv.PolyData = result["registered_template_model_surface"] | ||
| registered_file = output_dir / "registered_template.vtp" | ||
| registered_surface.save(str(registered_file)) | ||
| registered_coefficients = workflow.pca_coefficients | ||
| registered_coefficients_path = output_dir / "registered_coefficients.json" | ||
| with registered_coefficients_path.open(mode="w", encoding="utf-8") as f: | ||
| json.dump(registered_coefficients.tolist(), f) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
PCA registration unconditionally enabled even when pca_model is None — will raise ValueError.
use_pca_registration=True is hardcoded regardless of whether pca_model was actually loaded. pca_json is only conditionally loaded (line 83-85), and the docs' Recommended Run Order note says Tutorial 4 "can consume" the PCA model output, implying it's optional. But WorkflowFitStatisticalModelToPatient.set_use_pca_registration raises ValueError("When enabling PCA registration, pca_model must be provided.") whenever use_pca_registration=True and pca_model is None, so this tutorial will crash if Tutorial 3's pca_model.json hasn't been generated yet.
Compare with docs/tutorials.rst's Tutorial 4 "Inner API usage" snippet, which correctly guards this call: if pca_model is not None: workflow.set_use_pca_registration(True, pca_model=pca_model).
Additionally, even if this call didn't raise, registered_coefficients = workflow.pca_coefficients (line 140) would be None when PCA is skipped, and .tolist() (line 143) would then raise AttributeError.
🐛 Proposed fix
- workflow.set_use_pca_registration(
- use_pca_registration=True,
- pca_model=pca_model,
- use_surface=False,
- )
+ if pca_model is not None:
+ workflow.set_use_pca_registration(
+ use_pca_registration=True,
+ pca_model=pca_model,
+ use_surface=False,
+ )
# %%
# Workflow execution
workflow_results = workflow.process()
# %%
# Result saving
registered_coefficients = workflow.pca_coefficients
- registered_coefficients_path = output_dir / "registered_coefficients.json"
- with registered_coefficients_path.open(mode="w", encoding="utf-8") as f:
- json.dump(registered_coefficients.tolist(), f)
+ if registered_coefficients is not None:
+ registered_coefficients_path = output_dir / "registered_coefficients.json"
+ with registered_coefficients_path.open(mode="w", encoding="utf-8") as f:
+ json.dump(registered_coefficients.tolist(), f)Based on learnings and cross-file evidence: docs/tutorials.rst's Tutorial 4 usage snippet demonstrates the guarded pattern if pca_model is not None: workflow.set_use_pca_registration(...), which this tutorial script does not follow.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| workflow.set_use_pca_registration( | |
| use_pca_registration=True, | |
| pca_model=pca_model, | |
| use_surface=False, | |
| ) | |
| if pca_model is not None: | |
| workflow.set_use_pca_registration(True, pca_model=pca_model) | |
| # %% | |
| # Workflow execution | |
| result = workflow.run_workflow() | |
| workflow_results = workflow.process() | |
| # %% | |
| # Result saving | |
| registered_surface: pv.PolyData = result["registered_template_model_surface"] | |
| registered_file = output_dir / "registered_template.vtp" | |
| registered_surface.save(str(registered_file)) | |
| registered_coefficients = workflow.pca_coefficients | |
| registered_coefficients_path = output_dir / "registered_coefficients.json" | |
| with registered_coefficients_path.open(mode="w", encoding="utf-8") as f: | |
| json.dump(registered_coefficients.tolist(), f) | |
| if pca_model is not None: | |
| workflow.set_use_pca_registration( | |
| use_pca_registration=True, | |
| pca_model=pca_model, | |
| use_surface=False, | |
| ) | |
| # %% | |
| # Workflow execution | |
| workflow_results = workflow.process() | |
| # %% | |
| # Result saving | |
| registered_coefficients = workflow.pca_coefficients | |
| if registered_coefficients is not None: | |
| registered_coefficients_path = output_dir / "registered_coefficients.json" | |
| with registered_coefficients_path.open(mode="w", encoding="utf-8") as f: | |
| json.dump(registered_coefficients.tolist(), f) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tutorials/tutorial_04_fit_statistical_model_to_patient.py` around lines 128 -
143, Guard PCA registration in the workflow setup using the loaded pca_model:
call set_use_pca_registration only when pca_model is not None, while preserving
the existing enabled configuration when it is available. Update the
result-saving flow around workflow.pca_coefficients so it does not call tolist()
when PCA was skipped, and save registered coefficients only when they exist.
Non-contrast heart models from TotalSegmentator are incomplete, but this update makes registration to a statistical model possible.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
process()entry point and current segmentation terminology.