examples/example_carson.py

#!/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()