File size: 2,993 Bytes
2069947 |
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 |
import numpy as np
from scipy import ndimage
def extract_tumor_and_peritumoral(mask_volume, peritumoral_margin=2, patch_size=(16, 16, 16)):
"""
Extract tumor and peritumoral regions from a 3D annotation mask.
Flattens dilated mask into sequence of patches and creates position mask.
Parameters:
mask_volume: 3D numpy array (z, y, x) with tumor annotations (1 for tumor, 0 for background)
peritumoral_margin: Integer specifying the margin size (in voxels) for peritumoral region
patch_size: Tuple (z,y,x) specifying size of patches to use
Returns:
tumor_coords: List of coordinates (z, y, x) for tumor region
peritumoral_coords: List of coordinates (z, y, x) for peritumoral region
patch_mask: Binary mask indicating if patches contain tumor (1) or not (0)
"""
# Get tumor coordinates
tumor_coords = np.where(mask_volume == 1)
tumor_coords = list(zip(tumor_coords[0], tumor_coords[1], tumor_coords[2]))
# Create dilated mask for peritumoral region
dilated_mask = ndimage.binary_dilation(
mask_volume,
structure=np.ones((peritumoral_margin * 2 + 1, peritumoral_margin * 2 + 1, peritumoral_margin * 2 + 1)),
)
# Create patch position mask
z_steps = mask_volume.shape[0] // patch_size[0]
y_steps = mask_volume.shape[1] // patch_size[1]
x_steps = mask_volume.shape[2] // patch_size[2]
patch_mask = np.zeros((z_steps, y_steps, x_steps))
for z in range(z_steps):
for y in range(y_steps):
for x in range(x_steps):
patch = dilated_mask[
z * patch_size[0] : (z + 1) * patch_size[0],
y * patch_size[1] : (y + 1) * patch_size[1],
x * patch_size[2] : (x + 1) * patch_size[2],
]
if np.any(patch):
patch_mask[z, y, x] = 1
return tumor_coords, patch_mask.flatten()
# Example usage
def main():
# Create sample data for testing
volume_shape = (96, 96, 96)
mask_volume = np.zeros(volume_shape)
# Create a synthetic tumor mask in the middle
mask_volume[40:60, 40:60, 40:60] = 1
# Test parameters
patch_size = (16, 16, 16)
peritumoral_margin = 5
# Call function and get results
tumor_coords, patch_mask = extract_tumor_and_peritumoral(
mask_volume, peritumoral_margin=peritumoral_margin, patch_size=patch_size
)
# Print test results
print(f"Volume shape: {volume_shape}")
print(f"Tumor volume: {len(tumor_coords)}")
print(f"Number of total patches: {patch_mask.shape}")
print(f"Number of patches containing tumor/peritumoral region: {np.sum(patch_mask)}")
# Validate results
assert len(tumor_coords) > 0, "No tumor coordinates found"
assert len(patch_mask) == np.prod(np.array(volume_shape) // np.array(patch_size)), "Incorrect patch mask size"
return tumor_coords, patch_mask
if __name__ == "__main__":
tumor_coords, patch_mask = main()
|