Here it is:
def robustEndPointDetection(endpoint, segmentation, aff, n=5):
''' Relocates automatically detected endpoints to the center of mass of the closest component
inside a local region around the endpoint (defined by n).
Takes the endpoint position, converts it to voxel coordinates with the affine matrix, then defines a region
of (2 * n) ^ 3 voxels centered around the endpoint. Then components inside the local region are treated
as separate objects. The minimum distance from theese objects to the endpoint is computed, and from
these, the object with the smallest distance to the endpoint is chosen to compute the centroid, which
is converted back to RAS with the affine matrix.
Arguments:
- endpoint <np.array>: position of the endpoint in RAS coordinates.
- segmentation <np.array>: numpy array corresponding to the croppedVolumeNode.
- aff <np.array>: affine matrix corresponding ot he nifti file.
- n <int>: defines size of the region around the endpoint that is analyzed for this method.
Returns:
- newEndpoint <np.array>: new position of the endpoint.
'''
from skimage.measure import regionprops, label
from scipy import ndimage
# Compute RAS coordinates in voxel coordinates with affine matrix
R0, A0, S0 = np.round(np.matmul(np.linalg.inv(aff), np.append(endpoint, 1.0))[:3]).astype(int)
# Mask the segmentation (Only region of interest)
maskedSegmentation = segmentation[np.max([0, S0 - n]): np.min([segmentation.shape[0], S0 + n]),
np.max([0, A0 - n]): np.min([segmentation.shape[1], A0 + n]),
np.max([0, R0 - n]): np.min([segmentation.shape[2], R0 + n])]
# Divide into different connected components
labelMask = label(maskedSegmentation)
labels = np.sort(np.unique(labelMask))
labels = np.delete(labels, np.where([labels == 0]))
labelMaskOneHot = np.zeros([len(labels), labelMask.shape[0], labelMask.shape[1], labelMask.shape[2]], dtype=np.uint8)
for idx, label in enumerate(labels):
labelMaskOneHot[idx][labelMask == label] = 1
invertedLabelMaskOneHot = np.ones_like(labelMaskOneHot) - labelMaskOneHot
# Get distance transform for each and get only closest component
distanceLabels = np.empty_like(labels, dtype = np.float)
for idx in range(len(labels)):
distanceLabels[idx] = ndimage.distance_transform_edt(invertedLabelMaskOneHot[idx])[invertedLabelMaskOneHot.shape[1] // 2][invertedLabelMaskOneHot.shape[2] // 2][invertedLabelMaskOneHot.shape[3] // 2]
mask = np.zeros_like(segmentation)
mask[np.max([0, S0 - n]): np.min([segmentation.shape[0], S0 + n]),
np.max([0, A0 - n]): np.min([segmentation.shape[1], A0 + n]),
np.max([0, R0 - n]): np.min([segmentation.shape[2], R0 + n])] = labelMaskOneHot[np.argmin(distanceLabels)]
# Get the centroid of the foregroud region
properties = regionprops(mask.astype(np.int), mask.astype(np.int))
centerOfMass = np.array(properties[0].centroid)[[2, 1, 0]]
# Return the new position of the endpoint as RAS coordinates
return np.matmul(aff, np.append(centerOfMass, 1.0))[:3]
In my case, I get the affine matrix using the nibabel package, directly from the NIfTI file itself:
aff = nib.load("/path/to/nifti.nii.gz").affine
I initially tried this, but the smoothing filters were either too aggressive on the small vessels (these would vanish) or too soft for larger vessels.
I will investigate this further.