Make labeling more robust

[NathanMolinier]

Description

I noticed an issue with the labeling algorithm: here the vertebra C6 does not appear and C5 is used 2 times

step1 raw	step2 raw	output with tags

Currently for labeling we project all the center of mass of the segmentations (vertebrae and discs) onto the canal. But sometimes this projection may not be 100% reliable.

Idea

To solve this we could perform the labeling only using the segmentations of the discs (project only the discs onto the centerline). Then identify the vertebrae that are in between to label them individually.

Example

Once we have the labels for the discs, lets say C2-C3 and C3-C4 we just need to identify the vertebrae that is in between the two and it will automatically be C3. To identify the vertebra, we could just find the middle point in between the center of mass of C2-C3 and C3-C4 then find out to which vertebra this point belong or is the closest to its center of mass.

[yw7]

Another/additional strategy is straightening the segmentation by the canal centerline, then getting the center of mass for sorting by the z-axis:

maybe something like this can work before getting the center of mass here:

totalspineseg/totalspineseg/utils/iterative_label.py

Lines 812 to 813 in 3fa677e

	# Get the indices of the center of mass for each label
	labels_indices = np.array(ndi.center_of_mass(np.isin(mask_labeled, labels), mask_labeled, labels))

using canal_centerline_indices from:

totalspineseg/totalspineseg/utils/iterative_label.py

Lines 543 to 544 in 3fa677e

	# Get the canal centerline indices to use for sorting the discs and vertebrae based on the prjection on the canal centerline
	canal_centerline_indices = _get_canal_centerline_indices(seg_data, canal_labels + cord_labels)

import numpy as np
import nibabel as nib
from scipy.interpolate import RegularGridInterpolator

def _get_straighten_seg(seg_data, canal_centerline_indices):
    # Create grid coordinates using np.indices
    indices = np.indices(seg_data.shape)
    
    # Sort the indices by z-axis
    canal_centerline_indices = canal_centerline_indices[canal_centerline_indices[:, 2].argsort()]

    # Calculate the median x and y coordinates of the centerline
    median_x, median_y = np.median(canal_centerline_indices[:, :2], axis=0)
    
    # Initialize offset arrays with median values
    offsets_x = np.full(seg_data.shape[2], seg_data.shape[0] // 2 - median_x)
    offsets_y = np.full(seg_data.shape[2], seg_data.shape[1] // 2 - median_y)
    
    # Update offsets for z-slices with centerline points
    mask = np.isin(indices[2, 0, 0], canal_centerline_indices[:, 2])
    offsets_x[mask] = seg_data.shape[0] // 2 - canal_centerline_indices[:, 0]
    offsets_y[mask] = seg_data.shape[1] // 2 - canal_centerline_indices[:, 1]
    
    # Build the warping field for all slices
    warping_field = np.stack([
        indices[0] - offsets_x[np.newaxis, np.newaxis, :],
        indices[1] - offsets_y[np.newaxis, np.newaxis, :],
        indices[2]
    ], axis=-1)

    # Interpolate the image data using RegularGridInterpolator
    straightened_seg = RegularGridInterpolator(
        (indices[0][:, 0, 0], indices[1][0, :, 0], indices[2][0, 0, :]), seg_data,
        method='nearest', bounds_error=False, fill_value=0
    )(warping_field)
    
    return straightened_seg

Output:

[NathanMolinier]

I'm not sure straightening the cord would solve the issue. It would maybe just reduce the risks of swap between vertebrae.

In this case, the two vertebrae were assigned to the same label.

[yw7]

I'm not sure straightening the cord would solve the issue. It would maybe just reduce the risks of swap between vertebrae.

In this case, the two vertebrae were assigned to the same label.

In this example (which BTW already fixed in the current version in main branch) the problem is that C6 is swapped with disc C5-C6, then if you try to take the vertebrae between C4-C5 and C5-C6 you might get both (depending on how you check of course), so we can test if straightening the segmentation can solve this issue.

[NathanMolinier]

Personally i would not consider straightening at all. It is definitely too long and we can do better without it for sure.

the problem is that C6 is swapped with disc C5-C6

Not sure I understand, if discs can be labeled as vertebrae there is definitely something going wrong

[yw7]

Personally i would not consider straightening at all. It is definitely too long and we can do better without it for sure.

Personally, I would implement and compare both strategies in terms of accuracy and speed before being "sure".

Not sure I understand, if discs can be labeled as vertebrae there is definitely something going wrong

Yes, I wasn't clear enough. I meant that the Z-index of the center of mass projection in the canal centerline of C6 was above C5-C6.