diff --git a/docs/api/workflows.rst b/docs/api/workflows.rst index 41502c5..1fc6227 100644 --- a/docs/api/workflows.rst +++ b/docs/api/workflows.rst @@ -119,7 +119,8 @@ VTK to USD anatomy_type="heart", ) - output_path = workflow.process() + result = workflow.process() + usd_file = result["usd_file"] Statistical Shape Modeling ========================== diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 9c5e594..6d752f6 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -368,7 +368,8 @@ Inner API usage anatomy_type="heart", separate_by_connectivity=True, ) - usd_file = workflow.process() + result = workflow.process() + usd_file = result["usd_file"] For callers who need more control than the workflow wrapper offers (e.g. applying a colormap or per-label anatomical splitting), use diff --git a/src/physiotwin4d/cli/convert_vtk_to_usd.py b/src/physiotwin4d/cli/convert_vtk_to_usd.py index 56feaea..a6d201b 100644 --- a/src/physiotwin4d/cli/convert_vtk_to_usd.py +++ b/src/physiotwin4d/cli/convert_vtk_to_usd.py @@ -230,9 +230,9 @@ def main() -> int: return 1 try: - out_path = workflow.process() + result = workflow.process() print("\nConversion completed successfully.") - print(f"Output: {out_path}") + print(f"Output: {result['usd_file']}") return 0 except Exception as e: print(f"\nError during conversion: {e}") diff --git a/src/physiotwin4d/workflow_convert_vtk_to_usd.py b/src/physiotwin4d/workflow_convert_vtk_to_usd.py index a5314fb..c3346ba 100644 --- a/src/physiotwin4d/workflow_convert_vtk_to_usd.py +++ b/src/physiotwin4d/workflow_convert_vtk_to_usd.py @@ -9,7 +9,7 @@ import logging from pathlib import Path -from typing import Literal, Optional, Sequence, Union +from typing import Any, Literal, Optional, Sequence, Union import pyvista as pv import vtk @@ -106,12 +106,13 @@ def __init__( "separate_by_connectivity and separate_by_cell_type cannot both be True" ) - def process(self) -> str: + def process(self) -> dict[str, Any]: """ Run the full workflow: convert meshes to USD, then apply the chosen appearance. Returns: - Path to the created USD file (str). + Dict with the results of the workflow: + - "usd_file" (str): Path to the created USD file. """ self.log_section("VTK to USD conversion workflow") @@ -168,7 +169,7 @@ def process(self) -> str: self.log_warning( "No mesh prims found under /World/%s", self.usd_project_name ) - return str(output_usd) + return {"usd_file": str(output_usd)} # Static merge has no time samples; pass None so only default time is used appearance_time_codes = None if self.static_merge else time_codes @@ -223,4 +224,4 @@ def process(self) -> str: primvar = None # next mesh: auto-pick again self.log_info("Workflow complete: %s", output_usd) - return str(output_usd) + return {"usd_file": str(output_usd)} diff --git a/src/physiotwin4d/workflow_infer_physicsnemo.py b/src/physiotwin4d/workflow_infer_physicsnemo.py index 019311b..3e64374 100644 --- a/src/physiotwin4d/workflow_infer_physicsnemo.py +++ b/src/physiotwin4d/workflow_infer_physicsnemo.py @@ -382,6 +382,7 @@ def create_deformation_field( stage: float, reference_image: itk.Image, output_directory: Optional[Path] = None, + reference_surface: Optional[Path] = None, ) -> dict[str, Any]: """Rasterize the inferred deformation onto a reference image grid. @@ -392,6 +393,15 @@ def create_deformation_field( (renormalized) reference-surface normal of those vertices. Empty voxels are zero. + By default the reference (undeformed) surface is reconstructed from the + PCA coefficients in the model's own frame. When the subject's reference + surface lives in a different world frame than the model template (e.g. a + patient scan whose statistical-model fit applied a pose transform not + captured by the shape coefficients), pass ``reference_surface`` so the + displacements are binned at the patient-space positions that actually + align with ``reference_image``. The network displacements themselves + depend only on the coefficients and stage, not on the binning positions. + Args: shape_parameters: JSON file with the subject PCA coefficient vector. stage: Target RR-interval fraction for the deformation. @@ -399,14 +409,34 @@ def create_deformation_field( (size, spacing, origin, direction). output_directory: If given, the two images are written there as compressed ``.mha`` files. + reference_surface: Optional mesh (volume or surface) whose extracted + surface supplies the binning positions and normals, overriding + the PCA reconstruction. Must share the mean-shape surface + topology (same point count and ordering); its surface is + extracted with the same ``dataset_surface`` algorithm used for + the model template, so a mesh built from the same PCA template + keeps the correspondence. Returns: Dict with ``deformation_field`` and ``normal_image`` (ITK vector - images) and, when written, their paths. + images), ``deformed_surface`` (the stage surface as ``pv.PolyData``) + and, when written, their paths. """ coeffs = pnt.load_pca_coefficients(shape_parameters) - mean_mesh, pca_model = self._load_pca_assets() - ref_points = pnt.reconstruct_reference_points(mean_mesh, pca_model, coeffs) + if reference_surface is not None: + patient_surface = cast( + pv.DataSet, pv.read(str(reference_surface)) + ).extract_surface(algorithm="dataset_surface") + ref_points = np.asarray(patient_surface.points, dtype=np.float32) + n_expected = len(self._mean_shape_coords) + if ref_points.shape[0] != n_expected: + raise ValueError( + f"reference_surface has {ref_points.shape[0]} surface points, " + f"expected {n_expected} (mean-shape topology)." + ) + else: + mean_mesh, pca_model = self._load_pca_assets() + ref_points = pnt.reconstruct_reference_points(mean_mesh, pca_model, coeffs) disps = self._predict_displacements(coeffs, stage) # Reference (undeformed) surface normals. @@ -452,19 +482,28 @@ def create_deformation_field( ref_points.shape[0], ) + # Deformed (stage) surface: reference positions displaced by the network, + # keeping the mean-shape topology. + deformed_surface = self._mean_surface.copy(deep=True) + deformed_surface.points = (ref_points + disps).astype(np.float32) + result: dict[str, Any] = { "deformation_field": deformation_image, "normal_image": normal_image, + "deformed_surface": deformed_surface, } if output_directory is not None: out_dir = Path(output_directory) out_dir.mkdir(parents=True, exist_ok=True) field_path = out_dir / "deformation_field.mha" normal_path = out_dir / "surface_normal_field.mha" + surface_path = out_dir / "deformed_surface.vtp" itk.imwrite(deformation_image, str(field_path), compression=True) itk.imwrite(normal_image, str(normal_path), compression=True) + deformed_surface.save(str(surface_path)) result["deformation_field_file"] = field_path result["normal_image_file"] = normal_path + result["deformed_surface_file"] = surface_path return result @staticmethod diff --git a/tutorials/tutorial_02_lung_ct_to_vtk.py b/tutorials/tutorial_02_lung_ct_to_vtk.py new file mode 100644 index 0000000..6256420 --- /dev/null +++ b/tutorials/tutorial_02_lung_ct_to_vtk.py @@ -0,0 +1,156 @@ +""" +Tutorial 2: CT Segmentation to VTK Surfaces + +Purpose +------- +Segment one 3D CT frame into anatomical groups and save a combined VTK +surface file. The output can be inspected directly in PyVista or used as +input for Tutorial 3. + +Data Required +------------- +Full data: ``data/DirLab-4DCT/Case1Pack_T??.mha`` +Test data: ``data/test/DirLab-4DCT/Case1Pack_T??.mha`` +""" + +# Imports +from __future__ import annotations + +import logging +from pathlib import Path + +import itk +import pyvista as pv + +from physiotwin4d import ( + ContourTools, + SegmentChestTotalSegmentator, + TestTools, + WorkflowConvertImageToVTK, +) + +# Only run if this script is not imported as a module + +# nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a +# multiprocessing.Pool. On Windows the spawn start method re-imports this +# script in each child; without the __name__ == "__main__" guard around +# top-level work, that re-import fires the segmenter again and Python's +# spawn-cascade detector raises RuntimeError. +if __name__ == "__main__": + # Data directory specification + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + project_name = "tutorial_02_lung" + + output_dir = tutorials_dir / "output" / project_name + + # In addition to the combined surface file always saved below, also + # save one VTP per anatomy group (e.g. heart.vtp, lung.vtp) and/or one + # VTP per individual anatomical structure (e.g. left_ventricle.vtp). + save_group_surfaces = True + save_label_surfaces = True + + test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" + else: + data_dir = repo_root / "data" / "DirLab-4DCT" + + frame_files = sorted(data_dir.glob("Case1Pack_T??.mha")) + + log_level = logging.INFO + + segmentation_method = SegmentChestTotalSegmentator(log_level=log_level) + segmentation_method.set_has_academic_license(True) + + # Directory setup and data reading + output_dir.mkdir(parents=True, exist_ok=True) + + if not frame_files: + raise FileNotFoundError( + "DirLab-4DCT frame data not found. Checked:\n" + + f" - {data_dir}\n" + + "See data/README.md for download instructions." + ) + + ct_file = frame_files[0] + ct_image = itk.imread(str(ct_file)) + + # Workflow initialization + + workflow = WorkflowConvertImageToVTK( + segmentation_method=segmentation_method, + log_level=log_level, + ) + + # Workflow execution + # + # surface_target_reduction decimates each exported VTP surface. + result = workflow.process( + input_image=ct_image, + surface_target_reduction=0.5, + extract_label_surfaces=save_label_surfaces, + ) + + # Result saving + surface_file = Path( + ContourTools.save_combined_surface( + result["surfaces"], + str(output_dir), + prefix="patient", + ) + ) + if save_group_surfaces: + ContourTools.save_surfaces( + result["surfaces"], str(output_dir), prefix="patient" + ) + if save_label_surfaces: + ContourTools.save_surfaces( + result["label_surfaces"], str(output_dir), prefix="patient" + ) + labelmap_file = output_dir / "patient_labelmap.mha" + itk.imwrite(result["labelmap"], str(labelmap_file), compression=True) + + # Testing + tt = TestTools( + class_name=project_name, + results_dir=output_dir, + log_level=log_level, + ) + + screenshots: list[Path] = [] + screenshots.append( + tt.save_screenshot_image_slice( + ct_image, + f"{project_name}_segmentation_overlay.png", + axis=0, + slice_fraction=0.5, + colormap="gray", + vmin=-200, + vmax=600, + overlay_mask=result["labelmap"], + ) + ) + + surfaces = [ + surface for surface in result["surfaces"].values() if surface is not None + ] + if surfaces: + combined_surface = pv.merge(surfaces) if len(surfaces) > 1 else surfaces[0] + screenshots.append( + tt.save_screenshot_mesh( + combined_surface, + f"{project_name}_vtk_surfaces.png", + camera_position="iso", + color="lightblue", + opacity=0.85, + ) + ) + + tutorial_results = { + "result": result, + "surface_file": surface_file, + "labelmap_file": labelmap_file, + "screenshots": screenshots, + } diff --git a/tutorials/tutorial_03_heart_vtk_to_usd.py b/tutorials/tutorial_03_heart_vtk_to_usd.py index 09063a1..750b194 100644 --- a/tutorials/tutorial_03_heart_vtk_to_usd.py +++ b/tutorials/tutorial_03_heart_vtk_to_usd.py @@ -9,7 +9,6 @@ Data Required ------------- Preferred input: ``tutorials/output/tutorial_02_heart/patient_surfaces.vtp`` -Fallback input: any ``*.vtp`` under ``data`` or ``data/test`` """ # Imports @@ -17,7 +16,6 @@ import logging from pathlib import Path -from typing import Optional import pyvista as pv @@ -43,18 +41,12 @@ output_dir = tutorials_dir / "output" / "tutorial_03_heart" baselines_dir = repo_root / "tests" / "baselines" - # Preferred input: the combined surface saved by Tutorial 2. Leave vtk_file as - # None to auto-discover (Tutorial 2 output first, then any *.vtp under data_dir). - tutorial_02_surface = ( + project_name = "tutorial_02_heart" + + # Preferred input: the combined surface saved by Tutorial 2. + vtk_file = ( tutorials_dir / "output" / "tutorial_02_heart" / "patient_surfaces.vtp" ) - vtk_file: Optional[Path] = None - - test_mode = TestTools.running_as_test() - if test_mode: - data_dir = repo_root / "data" / "test" - else: - data_dir = repo_root / "data" log_level = logging.INFO @@ -62,24 +54,13 @@ output_dir.mkdir(parents=True, exist_ok=True) - if vtk_file is None and tutorial_02_surface.exists(): - vtk_file = tutorial_02_surface - if vtk_file is None: - vtk_candidates = sorted(data_dir.rglob("*.vtp")) - if not vtk_candidates: - raise FileNotFoundError( - "No VTK surface file found. Run Tutorial 2 first or place a " - f"*.vtp file under {data_dir}." - ) - vtk_file = vtk_candidates[0] - mesh = pv.read(str(vtk_file)) # Workflow initialization workflow = WorkflowConvertVTKToUSD( input_meshes=[mesh], - usd_project_name="surfaces", + usd_project_name=project_name, output_directory=output_dir, appearance="anatomy", anatomy_type="heart", @@ -88,7 +69,7 @@ ) # Workflow execution - usd_file = workflow.process() + results = workflow.process() # Testing tt = TestTools( @@ -98,12 +79,11 @@ log_level=log_level, ) - screenshots: list[Path] = [] - screenshots.append( + screenshots = [ tt.save_screenshot_openusd( - usd_file, - "usd_mesh_rendering.png", + results["usd_file"], + f"{project_name}_usd_mesh_rendering.png", ) - ) + ] - tutorial_results = {"usd_file": usd_file, "screenshots": screenshots} + tutorial_results = {"usd_file": results["usd_file"], "screenshots": screenshots} \ No newline at end of file diff --git a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py index 7354759..83a5993 100644 --- a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py @@ -45,9 +45,9 @@ repo_root = Path(__file__).resolve().parent.parent tutorials_dir = Path(__file__).resolve().parent - class_name = "tutorial_05_heart_to_lung_fit_statistical_model_to_patient" + project_name = "tutorial_05_heart_to_lung" - output_dir = tutorials_dir / "output" / "tutorial_05_heart_to_lung" + output_dir = tutorials_dir / "output" / project_name baselines_dir = repo_root / "tests" / "baselines" # PCA model + mean surface produced by Tutorial 4. @@ -55,6 +55,9 @@ pca_mean_file = ( tutorials_dir / "output" / "tutorial_04_heart" / "pca_mean_surface.vtp" ) + # BYOD example: + # pca_mean_file = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + # pca_json = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_model.json") test_mode = TestTools.running_as_test() if test_mode: @@ -94,18 +97,25 @@ "See data/README.md for download instructions." ) patient_image = itk.imread(str(patient_image_file)) - itk.imwrite(patient_image, output_dir / "patient_image.nii.gz") - segmentation_result = segmentation_method.segment(patient_image) - patient_labelmap = segmentation_result["labelmap"] - itk.imwrite(patient_labelmap, output_dir / "patient_labelmap.nii.gz") + if not (output_dir / f"{project_name}_patient_image.nii.gz").exists(): + itk.imwrite(patient_image, output_dir / f"{project_name}_patient_image.nii.gz") - heart_labelmap = segmentation_result["heart"] - itk.imwrite(heart_labelmap, output_dir / "heart_labelmap.nii.gz") + segmentation_result = segmentation_method.segment(patient_image) + patient_labelmap = segmentation_result["labelmap"] + itk.imwrite(patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz") - contour_tools = ContourTools() - heart_surface = contour_tools.extract_contours(labelmap_image=heart_labelmap) - heart_surface.save(output_dir / "heart_surface.vtp") + heart_labelmap = segmentation_result["heart"] + itk.imwrite(heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz") + + contour_tools = ContourTools() + heart_surface = contour_tools.extract_contours(labelmap_image=heart_labelmap) + heart_surface.save(output_dir / f"{project_name}_heart_surface.vtp") + + else: + patient_labelmap = itk.imread(output_dir / f"{project_name}_patient_labelmap.nii.gz") + heart_labelmap = itk.imread(output_dir / f"{project_name}_heart_labelmap.nii.gz") + heart_surface = pv.read(output_dir / f"{project_name}_heart_surface.vtp") # Workflow initialization @@ -131,25 +141,25 @@ # Result saving registered_coefficients = workflow.pca_coefficients if registered_coefficients is not None: - registered_coefficients_path = output_dir / "registered_coefficients.json" + registered_coefficients_path = output_dir / f"{project_name}_registered_coefficients.json" with registered_coefficients_path.open(mode="w", encoding="utf-8") as f: json.dump(registered_coefficients.tolist(), f) template_mesh = workflow.pca_template_model - template_mesh.save(str(output_dir / "template_mesh.vtp")) + template_mesh.save(str(output_dir / f"{project_name}_template_mesh.vtu")) template_surface = workflow.pca_template_model_surface - template_surface.save(str(output_dir / "template_surface.vtp")) + template_surface.save(str(output_dir / f"{project_name}_template_surface.vtp")) registered_mesh = workflow_results["registered_template_model"] - registered_mesh.save(str(output_dir / "template_mesh_registered.vtp")) + registered_mesh.save(str(output_dir / f"{project_name}_template_mesh_registered.vtu")) registered_surface = workflow_results["registered_template_model_surface"] - registered_surface.save(str(output_dir / "template_surface_registered.vtp")) + registered_surface.save(str(output_dir / f"{project_name}_template_surface_registered.vtp")) # Testing TestTools( - class_name=class_name, + class_name=project_name, results_dir=output_dir, baselines_dir=baselines_dir, log_level=log_level, @@ -162,7 +172,7 @@ screenshots: list[Path] = [] - before_path = output_dir / "model_before_registration.png" + before_path = output_dir / f"{project_name}_model_before_registration.png" plotter = pv.Plotter(off_screen=True, window_size=[800, 600]) plotter.add_mesh(pca_mean, color="dodgerblue", opacity=0.6) plotter.add_mesh(heart_surface, color="tomato", opacity=0.6) @@ -171,7 +181,7 @@ plotter.close() screenshots.append(before_path) - after_path = output_dir / "model_after_registration.png" + after_path = output_dir / f"{project_name}_model_after_registration.png" plotter = pv.Plotter(off_screen=True, window_size=[800, 600]) plotter.add_mesh(registered_surface, color="limegreen", opacity=0.7) plotter.add_mesh(heart_surface, color="tomato", opacity=0.4) diff --git a/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py index ec664fe..a6276e5 100644 --- a/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py +++ b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py @@ -90,7 +90,7 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" MANIFESTS_DIR = OUTPUT_DIR / "manifests_mgn" - RESUME_FROM = str(OUTPUT_DIR / "mgn_stage_model_epoch_00100.pt") + RESUME_FROM = str(TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn_2" / "mgn_stage_model_epoch_00200.pt") EPOCHS = 1500 BATCH_SIZE_GRAPHS = 4 # mini-batch measured in (subject, phase) graphs