#!/usr/bin/env python """ Example usage of CarSONSegmentor for cardiac MRI segmentation. This script demonstrates how to use the ONNX-based CarSON segmentor to segment cardiac structures from MRI images. """ import numpy as np import matplotlib.pyplot as plt from pathlib import Path import sys # Add parent directory to path for imports sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) from cardiac_toolkit.segmentation import CarSONSegmentor def example_basic_usage(): """Example 1: Basic 2D segmentation.""" print("=" * 80) print("Example 1: Basic 2D Segmentation") print("=" * 80) # Initialize segmentor # Default model path (update with your actual path) model_path = "/home/jx332/project/code/2025-05-DeepStrain/pretrained_models/carson_Jan2021.onnx" if not Path(model_path).exists(): print(f"\n⚠ Model not found at {model_path}") print("Please update the model path in this script.") return None, None segmentor = CarSONSegmentor(model_path=model_path) # Create dummy image (replace with your own image loading) # In practice: image = np.load('your_cardiac_mri.npy') image = np.random.rand(256, 256).astype(np.float32) * 1000 print(f"\nInput image shape: {image.shape}") print(f"Input image dtype: {image.dtype}") print(f"Input image range: [{image.min():.2f}, {image.max():.2f}]") # Perform segmentation print("\nPerforming segmentation...") mask = segmentor.segment(image) print(f"\nOutput mask shape: {mask.shape}") print(f"Output mask dtype: {mask.dtype}") print(f"Unique labels: {np.unique(mask)}") # Count pixels per class print("\nPixel counts per class:") for idx, name in segmentor.get_class_names().items(): count = (mask == idx).sum() print(f" {name} (class {idx}): {count} pixels") return image, mask def example_with_probabilities(): """Example 2: Getting probability maps.""" print("\n" + "=" * 80) print("Example 2: Probability Maps") print("=" * 80) # Initialize segmentor model_path = "/home/jx332/project/code/2025-05-DeepStrain/pretrained_models/carson_Jan2021.onnx" if not Path(model_path).exists(): print(f"\n⚠ Model not found at {model_path}") return None, None segmentor = CarSONSegmentor(model_path=model_path) # Create dummy image image = np.random.rand(256, 256).astype(np.float32) * 1000 print(f"\nInput image shape: {image.shape}") # Get probability maps print("\nGetting probability maps...") probs = segmentor.predict(image, return_probs=True) print(f"\nOutput probabilities shape: {probs.shape}") print(f"Probability range: [{probs.min():.3f}, {probs.max():.3f}]") # Check that probabilities sum to 1 (in center region where predictions exist) h, w = image.shape center_h, center_w = h // 2, w // 2 center_probs = probs[center_h-50:center_h+50, center_w-50:center_w+50, :] prob_sum = center_probs.sum(axis=-1) print(f"\nProbability sum per pixel (center region):") print(f" Mean: {prob_sum.mean():.6f}") print(f" Std: {prob_sum.std():.6f}") return image, probs def example_3d_volume(): """Example 3: 3D volume segmentation.""" print("\n" + "=" * 80) print("Example 3: 3D Volume Segmentation") print("=" * 80) # Initialize segmentor model_path = "/home/jx332/project/code/2025-05-DeepStrain/pretrained_models/carson_Jan2021.onnx" if not Path(model_path).exists(): print(f"\n⚠ Model not found at {model_path}") return None, None segmentor = CarSONSegmentor(model_path=model_path, batch_size=16) # Create dummy 3D volume volume = np.random.rand(256, 256, 10).astype(np.float32) * 1000 print(f"\nInput volume shape: {volume.shape}") # Segment entire volume print("\nSegmenting 3D volume...") mask_3d = segmentor.segment(volume) print(f"\nOutput mask shape: {mask_3d.shape}") print(f"Unique labels: {np.unique(mask_3d)}") # Count pixels per class across all slices print("\nPixel counts per class (all slices):") for idx, name in segmentor.get_class_names().items(): count = (mask_3d == idx).sum() print(f" {name} (class {idx}): {count} voxels") return volume, mask_3d def example_visualization(): """Example 4: Visualizing results.""" print("\n" + "=" * 80) print("Example 4: Visualization") print("=" * 80) # Initialize segmentor model_path = "/home/jx332/project/code/2025-05-DeepStrain/pretrained_models/carson_Jan2021.onnx" if not Path(model_path).exists(): print(f"\n⚠ Model not found at {model_path}") return segmentor = CarSONSegmentor(model_path=model_path) # Create dummy image image = np.random.rand(256, 256).astype(np.float32) * 1000 # Perform segmentation mask = segmentor.segment(image) probs = segmentor.predict(image, return_probs=True) # Create visualization print("\nCreating visualization...") fig, axes = plt.subplots(2, 3, figsize=(15, 10)) # Original image axes[0, 0].imshow(image, cmap='gray') axes[0, 0].set_title('Original Image') axes[0, 0].axis('off') # Segmentation mask axes[0, 1].imshow(mask, cmap='jet', vmin=0, vmax=3) axes[0, 1].set_title('Segmentation\n(0=BG, 1=LV, 2=Myo, 3=RV)') axes[0, 1].axis('off') # Overlay axes[0, 2].imshow(image, cmap='gray') axes[0, 2].imshow(mask, cmap='jet', alpha=0.4, vmin=0, vmax=3) axes[0, 2].set_title('Overlay') axes[0, 2].axis('off') # Class probabilities class_names = ['Background', 'LV Cavity', 'Myocardium'] for i in range(3): ax = axes[1, i] im = ax.imshow(probs[:, :, i+1], cmap='hot', vmin=0, vmax=1) ax.set_title(f'{class_names[i]} Probability') ax.axis('off') plt.colorbar(im, ax=ax, fraction=0.046) plt.tight_layout() # Save figure output_path = Path(__file__).parent / 'carson_segmentation_result.png' plt.savefig(output_path, dpi=150, bbox_inches='tight') print(f"\nVisualization saved to: {output_path}") plt.close() def main(): """Run all examples.""" print("\n" + "*" * 80) print("CarSONSegmentor Examples") print("*" * 80) try: # Example 1: Basic usage example_basic_usage() # Example 2: Probability maps example_with_probabilities() # Example 3: 3D volume example_3d_volume() # Example 4: Visualization example_visualization() print("\n" + "=" * 80) print("All examples completed successfully!") print("=" * 80) print() except Exception as e: print(f"\n\nError running examples: {e}") import traceback traceback.print_exc() if __name__ == '__main__': main()