Upload tumor_mask.py

tumor_mask.py

@@ -0,0 +1,84 @@
+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()