File size: 9,059 Bytes
d5d23f9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
#!/usr/bin/env python
"""
Example usage of JDLSegmentor for myocardium and scar segmentation.

This script demonstrates how to use the ONNX-based JDL segmentor
to segment myocardium and scar tissue from cardiac 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 JDLSegmentor


def example_basic_usage():
    """Example 1: Basic 2D segmentation."""
    print("=" * 80)
    print("Example 1: Basic 2D Segmentation")
    print("=" * 80)

    # Initialize segmentor
    # Default model paths (update with your actual paths)
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        print(f"Myocardium: {myocardium_model_path}")
        print(f"Scar: {scar_model_path}")
        print("Please update the model paths in this script.")
        return None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_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...")
    myo_mask, scar_mask = segmentor.segment(image)

    print(f"\nMyocardium mask shape: {myo_mask.shape}")
    print(f"Myocardium pixels: {np.sum(myo_mask)}")
    print(f"Scar mask shape: {scar_mask.shape}")
    print(f"Scar pixels: {np.sum(scar_mask)}")

    return image, myo_mask, scar_mask


def example_with_bbox():
    """Example 2: Segmentation with bounding box."""
    print("\n" + "=" * 80)
    print("Example 2: Segmentation with Bounding Box")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return None, None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy image
    image = np.random.rand(512, 512).astype(np.float32) * 1000

    # Define a bounding box around the heart region
    bbox = {
        'x1': 150,
        'y1': 150,
        'x2': 350,
        'y2': 350
    }

    print(f"\nInput image shape: {image.shape}")
    print(f"Bounding box: {bbox}")

    # Perform segmentation with bbox
    print("\nPerforming segmentation with bounding box...")
    myo_mask, scar_mask = segmentor.segment(
        image,
        bbox=bbox,
        bbox_scale_factor=1.1  # Expand bbox by 10%
    )

    print(f"\nMyocardium mask shape: {myo_mask.shape}")
    print(f"Myocardium pixels: {np.sum(myo_mask)}")
    print(f"Scar mask shape: {scar_mask.shape}")
    print(f"Scar pixels: {np.sum(scar_mask)}")

    return image, myo_mask, scar_mask, bbox


def example_batch_processing():
    """Example 3: Batch processing of multiple slices."""
    print("\n" + "=" * 80)
    print("Example 3: Batch Processing")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy batch of images (e.g., multiple slices)
    num_slices = 5
    images = [np.random.rand(256, 256).astype(np.float32) * 1000 for _ in range(num_slices)]

    print(f"\nNumber of slices: {num_slices}")
    print(f"Each slice shape: {images[0].shape}")

    # Process batch
    print("\nProcessing batch...")
    myo_masks, scar_masks = segmentor.segment_batch(images)

    print(f"\nProcessed {len(myo_masks)} slices")
    for i, (myo_mask, scar_mask) in enumerate(zip(myo_masks, scar_masks)):
        print(f"  Slice {i}: Myo pixels = {np.sum(myo_mask)}, Scar pixels = {np.sum(scar_mask)}")

    return images, myo_masks, scar_masks


def example_3d_volume():
    """Example 4: 3D volume segmentation."""
    print("\n" + "=" * 80)
    print("Example 4: 3D Volume Segmentation")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy 3D volume
    volume = np.random.rand(256, 256, 8).astype(np.float32) * 1000

    print(f"\nInput volume shape: {volume.shape}")

    # Segment 3D volume
    print("\nSegmenting 3D volume...")
    myo_mask_3d, scar_mask_3d = segmentor.segment_3d(
        volume,
        keep_largest_component='per_slice'  # or 'whole_volume'
    )

    print(f"\nMyocardium mask shape: {myo_mask_3d.shape}")
    print(f"Myocardium voxels: {np.sum(myo_mask_3d)}")
    print(f"Scar mask shape: {scar_mask_3d.shape}")
    print(f"Scar voxels: {np.sum(scar_mask_3d)}")

    return volume, myo_mask_3d, scar_mask_3d


def example_visualization():
    """Example 5: Visualizing results."""
    print("\n" + "=" * 80)
    print("Example 5: Visualization")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy image
    image = np.random.rand(256, 256).astype(np.float32) * 1000

    # Perform segmentation
    print("\nPerforming segmentation...")
    myo_mask, scar_mask = segmentor.segment(image)

    # Create visualization
    print("\nCreating visualization...")
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))

    # Original image
    axes[0].imshow(image, cmap='gray')
    axes[0].set_title('Original Image')
    axes[0].axis('off')

    # Myocardium mask
    axes[1].imshow(myo_mask, cmap='gray')
    axes[1].set_title('Myocardium Mask')
    axes[1].axis('off')

    # Scar mask
    axes[2].imshow(scar_mask, cmap='gray')
    axes[2].set_title('Scar Mask')
    axes[2].axis('off')

    # Overlay
    axes[3].imshow(image, cmap='gray')
    axes[3].imshow(myo_mask, cmap='Reds', alpha=0.3)
    axes[3].imshow(scar_mask, cmap='Blues', alpha=0.3)
    axes[3].set_title('Overlay\n(Red=Myo, Blue=Scar)')
    axes[3].axis('off')

    plt.tight_layout()

    # Save figure
    output_path = Path(__file__).parent / 'jdl_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("JDLSegmentor Examples")
    print("*" * 80)

    try:
        # Example 1: Basic usage
        example_basic_usage()

        # Example 2: With bounding box
        example_with_bbox()

        # Example 3: Batch processing
        example_batch_processing()

        # Example 4: 3D volume
        example_3d_volume()

        # Example 5: 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()