Meta-learning for cardiac MRI segmentation
Methods and systems are described for image segmentation. A machine learning model is applied to a set of images to generate results. The results may be obtained as a probability map for each image in the set of images. The model may be trained by accessing a set of labeled images, each image associated with a label indicating a location of a feature within a respective image. An initial set of parameters is accessed. An encoder is initialized with the initial set of parameters. The encoder is applied to the set of labeled images to generate a prediction of a feature location within each image. The initial set of parameters are updated based on the predictions and the label associated with the labeled images. The updated set of parameters and an additional set of parameters generated using a set of unlabeled images are aggregated.
1. A computer implemented method of segmenting a set of images comprising:
applying a machine learning model to the set of images to generate results, wherein the machine learning model was trained by:
initializing a first encoder with a first set of parameters,
applying the first encoder to a set of unlabeled images to generate a set of representations of the unlabeled images,
updating the first set of parameters based on the set of representations, initializing a second encoder with a second set of parameters,
applying the second encoder to a set of labeled images to generate a set of predictions of a location of a feature within each image in the set of labeled images, wherein each respective label is associated with the location of the feature in the respective labeled image,
updating the second set of parameters based on the set of predictions, aggregating the updated first set of parameters and the updated second set of parameters into a set of aggregated parameters, and
updating the first set of parameters based on the set of aggregated parameters; and
obtaining the results of the machine learning model as a probability map for each image in the set of images.
2. The method of claim 1 , wherein the second set of parameters is the updated first set of parameters.
3. The method of claim 1 , wherein the first encoder is optimized according to a self-supervised loss function.
4. The method of claim 3 , wherein optimizing the first encoder comprises:
augmenting an instance of an image in the set of unlabeled imaged with one or more transformation operators to generate a plurality of augmented instances;
encoding the plurality of augmented instances to generate representations of the plurality of augmented instances; and
applying a similarity metric to the representations of the plurality of augmented instances.
5. The method of claim 1 , wherein the second encoder is optimized according to a cross entropy loss function.
6. The method of claim 5 , wherein optimizing the second encoder comprises:
mapping the set of labeled images into D-dimensional representations with the second encoder; and
mapping the representations into segmentation maps consisting of a set of classes with a decoder.
7. The method of claim 1 , wherein the first encoder is associated with a first task and wherein the second encoder is associated with a second task.
8. The method of claim 7 , wherein aggregating the updated first set of parameters and the updated second set of parameters comprises generating new parameters based on an update step that is defined by:
θ
=
θ
-
α
∑
t
=
1
T
(
θ
-
θ
t
)
,
wherein θ are initialization parameters, α is a learning rate, t is a task, and θ t are task-specific parameters.
9. The method of claim 7 , wherein aggregating the updated first set of parameters and the updated second set of parameters comprises generating new parameters based on an update step that is defined by:
θ=θ−α∇ E T [d(θ,M)], wherein θ are initialization parameters, α is a learning rate, and M is a manifold associated with each task, t, in a set of T tasks.
10. The method of claim 1 , wherein the set of unlabeled images and the set of labeled images are each associated with at least one of a first species, a first chamber, or a first modality, the method further comprising testing the machine learning model with a set of testing images, wherein the set of testing images are associated with at least one of a second species, a second chamber, or a second modality.
11. A non-transitory computer-readable storage medium having instructions encoded thereon that, when executed by one or more processors, cause the one or more processors to:
apply a machine learning model to a set of images to generate results, wherein the machine learning model was trained by:
initializing a first encoder with a first set of parameters,
applying the first encoder to a set of unlabeled images to generate a set of representations of the unlabeled images,
updating the first set of parameters based on the set of representations,
initializing a second encoder with a second set of parameters,
applying the second encoder to a set of labeled images to generate a set of predictions of a location of a feature within each image in the set of labeled images, wherein each respective label is associated with the location of the feature in the respective labeled image,
updating the second set of parameters based on the set of predictions,
aggregating the updated first set of parameters and the updated second set of parameters into a set of aggregated parameters, and
updating the first set of parameters based on the set of aggregated parameters; and
obtain the results of the machine learning model as a probability map for each image in the set of images.
12. The non-transitory computer-readable storage medium of claim 11 , wherein the second set of parameters is the updated first set of parameters.
13. The non-transitory computer-readable storage medium of claim 11 , wherein the first encoder is optimized according to a self-supervised loss function.
14. The non-transitory computer-readable storage medium of claim 13 , wherein optimizing the first encoder comprises:
augmenting an instance of an image in the set of unlabeled imaged with one or more transformation operators to generate a plurality of augmented instances;
encoding the plurality of augmented instances to generate representations of the plurality of augmented instances; and
applying a similarity metric to the representations of the plurality of augmented instances.
15. The non-transitory computer-readable storage medium of claim 11 , wherein the second encoder is optimized according to a cross entropy loss function.
16. The non-transitory computer-readable storage medium of claim 15 , wherein optimizing the second encoder comprises:
mapping the set of labeled images into D-dimensional representations with the second encoder; and
mapping the representations into segmentation maps consisting of a set of classes with a decoder.
17. The non-transitory computer-readable storage medium of claim 11 , wherein the first encoder is associated with a first task and wherein the second encoder is associated with a second task.
18. The non-transitory computer-readable storage medium of claim 17 , wherein aggregating the updated first set of parameters and the updated second set of parameters comprises generating new parameters based on an update step that is defined by:
θ
=
θ
-
α
∑
t
=
1
T
(
θ
-
θ
t
)
,
wherein θ are initialization parameters, α is a learning rate, t is a task, and θ t are task-specific parameters.
19. The non-transitory computer-readable storage medium of claim 17 , wherein aggregating the updated first set of parameters and the updated second set of parameters comprises generating new parameters based on an update step that is defined by:
θ=θ−α∇ T [d(θ,M)], wherein θ are initialization parameters, α is a learning rate, and M is a manifold associated with each task, t, in a set of T tasks.
20. The non-transitory computer-readable storage medium of claim 11 , wherein the set of unlabeled images and the set of labeled images are each associated with at least one of a first species, a first chamber, or a first modality, and
wherein the instructions further cause the one or more processors to test the machine learning model with a set of testing images, wherein the set of testing images are associated with at least one of a second species, a second chamber, or a second modality.