Role of mask_a versus label_a in forward

[aylward]

uniGradICON/src/unigradicon/__init__.py

Lines 43 to 49 in a96ec53

	if self.dice_loss_weight > 0.0:
	assert mask_A is not None and mask_B is not None, "mask_A and mask_B must be provided when dice_loss_weight>0"
	unique_A = torch.unique(mask_A.long())
	unique_B = torch.unique(mask_B.long())
	common_labels = unique_A[torch.isin(unique_A, unique_B)]
	common_labels = common_labels[common_labels != 0] # exclude background
	num_classes = len(common_labels)

I am uncertain of what form is expected for mask_a / mask_b and label_a / label_b. I assume that label_a/b would be a labelmap with multiple classes and used in computing dice scores. And that mask_a/b would be a binary image that is used to mask the image loss computation. However, as shown in the code above, the num_classes is computed from mask, and mask_a/b must be specified if dice weight > 0. Seems like label_a/b should be used if dice weight > 0 and label_a/b should be used in computing the num_classes.

[aylward]

@basardemir - @marcniethammer thought you might have some insights on this. I am using your fine-tune branch...and fine tuning using a labelmap and image seems to have worked well (but I am having trouble testing it).

[basardemir]

Thank you for raising this issue. I have been working on this branch for a while and have not merged it into main yet, but I will do so soon.

I agree that the variable naming is confusing, and I should come up with better names. In short, the label_a and label_b parameters are there for the "image similarity loss randomization" feature in multigradicon training . We used this notation in Equation 4 of the multigradicon paper. Referring to the paper:

Formally, we assume that our dataset (D = {P_i = {I_i^1, I_i^2, \ldots, I_i^m}}_{i \in [n]}) consists of (m) scans in different modalities for each patient (P_i). We first sample a patient pair ((P_A, P_B)). We then uniformly sample a source and target image pair ((I^A, I^B)) with (I^A \in P_A), (I^B \in P_B), and a source and target image pair ((I_L^A, I_L^B)) with (I_L^A \in P_A), (I_L^B \in P_B) for similarity loss computations. The pair ((I^A, I^B)) is used for transformation map prediction, and ((I_L^A, I_L^B)) is used for similarity loss calculation.

So, if you train a model by passing segmentations to the label parameters, that means you input images to the model but calculate the similarity loss on their segmentations. In that case, there is no loss term that incorporates the image intensities.

On the other hand, mask_a and mask_b are there for Dice losses. Initially, we used these terms for masking out the region of interest with provided masks, which we generally apply in lung datasets. Later, in order to preserve backward compatibility, I kept the same naming but also overloaded them with Dice loss calculations. If you use these parameters, the training loss will be calculated using both image intensities and Dice loss on the segmentations. I believe this was your intention.

I hope this clarifies the distinction, but I would like to hear your thoughts on it.

[aylward]

My preference is to not have the mask serve as both the dice loss and the roi.

For the ROI, I often (e.g., for cardiac CT) use a dilated segmentation. Dilation (binary) allows the registration to align the internals structure and the edges of an abject without getting significant interference from the adjacent structures. Without dilation, only internal structure is used.

However, when computing dice scores on segmentations, I want to consider class labels (not binary) and have the edges of the segmentation align with the edges of the objects (no dilation).

One option is to assign the dilated component of the object of interest to a new bogus class and pass that as the mask and leave label_a/b as None. But then it will calculate the dice using the bogus class boundary, and I don't want that.

I suggest adding a third image for computing dice scores, Segmentation_a/b. Mask_a/b will be strictly for ROI. Label_a/b will be unchanged (replacing image match). Segmentation_a/b will be for computing dice scores. This will maintain backward compatibility and be intuitive. This is also consistent with the finetuning json file that has a segmentation tagged image that is used for dice score computing.

[marcniethammer]

I agree with @aylward that it would be better to not overload the meanings of these terms.

We do want to be able to do loss function masking. But we also want to be able to use either the image or the labelmap or both as part of the similarity measure. The way I understand the current implementation (based on your description above) is that currently one can either have loss function masking or use a segmentation with a Dice loss but not both. And somehow the behavior is context dependent, which I think can get confusing.

[basardemir]

Hello @aylward,

I have updated the branch to clarify and separate the roles of masks and segmentations.

The CLI now supports separate arguments for each use case:

--fixed_mask, --moving_mask: binary ROI masks used for input masking and/or loss-function masking.
--fixed_segmentation, --moving_segmentation: integer label maps used for Dice loss when --dice_loss_weight > 0.

In other words, masks are no longer used to determine Dice classes. Dice loss now uses the provided segmentations, while masks are reserved for masking the input image or the similarity loss.

I also updated the README with examples for input masking, loss-function masking, Dice loss with segmentations, and using masks and segmentations together:

https://github.com/uncbiag/uniGradICON/tree/feat-add-finetuning

For backward compatibility, older commands that passed segmentations as masks are still accepted when Dice loss is not enabled, but this path is deprecated. Going forward, --fixed_mask / --moving_mask should be used for masks, and --fixed_segmentation / --moving_segmentation should be used for Dice-loss label maps.

We are looking forward to hearing your opinions and feedback after you have had a chance to try it.

[aylward]

Outstanding! Will try asap!

[aylward]

A few comments:

• In the generated file finetune.py, lines 199-204:

def _apply_roi_masking(batch: Dict[str, torch.Tensor]) -> None:
"""Zero out background under the ROI mask. Applied to both training and
validation batches so the network sees the same input distribution in
both phases."""
batch[PairKeys.IMAGE_A] = batch[PairKeys.IMAGE_A] * (batch[PairKeys.MASK_A] > 0).float()
batch[PairKeys.IMAGE_B] = batch[PairKeys.IMAGE_B] * (batch[PairKeys.MASK_B] > 0).float()

This doesn't "zero out the background under the ROI mask". The formula '''img * mask>0''' zeros the image where the mask is 0. So, "keeps everything under the mask" or "zeros out the image where the mask is 0" would be a better comment/description.

• I have tried multiple fine-tuning combinations, with only slight benefits resulting. I am registering gated (time-series) heart CT images. Since this is within patient, but the heart slides relative to adjacent structures (discontinuous at boundaries). Do you have suggestions on parameters? I have 29 patients with 5 to 20 time-point CT scans per patient for a total of 384 images. I've been running 250 epochs, using images, segementations (of regions of the heart) and a slightly dilated segmentation of the entire heart as the mask.

• I'm using both roi_masking and loss_function_masking, but I question the need for both options. Problem: If you don't mask the loss function but your do apply an roi_mask, the boundary of the roi_mask becomes a strong driver of the registration (it will try to make the ROI boundaries align. This only makes sense if boundary segmentation is perfect, which seems unlikely for any task. So, two questions (1) when would you want roi_masking without loss_function_masking? (2) what is the utility of roi_masking if you're doing loss function masking? I would anticipate no change in result if you use loss_function_masking without roi_masking. Proposal: Given these questions, is only a single "masking" option required, and it applies loss_function_masking, and roi_masking can be removed?

  • actually, the loss_function_mask is only uses mask_b (* lncc)...so yeah, roi_masking does impact the results of loss_function_masking...in an odd way. The image_A*roi_mask>0 will perhaps result in 0's warping into the space of mask_b, that is, the roi_mask boundary of image_A*roi_mask>0 will drive lncc computation, even if loss_function_mask is enabled. So, the use of only mask_B in loss function masking convinces me even more that loss_function_masking should be the only option (not roi_masking) and only a mask_b needs to be provided (mask_A is unused) - unless you can think of a use case...