Learning from partially labeled datasets for medical imaging analysis
Systems and methods for performing a medical imaging analysis task are provided. One or more input medical images are received. A medical imaging analysis task is performed based on the one or more input medical images using a machine learning based model. Results of the medical imaging analysis task are output. The machine learning based model is trained by receiving one or more training medical images and ground truth labels identifying one or more anatomical objects in the one or more training medical images, training the machine learning based model for performing the medical imaging analysis task based on the one or more training medical images and the ground truth labels, generating pseudo labels identifying the one or more anatomical objects in the one or more training medical images using the trained machine learning based model, and fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the generated pseudo labels.
1 . A computer-implemented method comprising:
receiving one or more input medical images;
performing a medical imaging analysis task based on the one or more input medical images using a machine learning based model; and
outputting results of the medical imaging analysis task,
wherein the machine learning based model is trained by:
receiving one or more training medical images and ground truth labels identifying one or more anatomical objects in the one or more training medical images,
training the machine learning based model for performing the medical imaging analysis task based on the one or more training medical images and the ground truth labels,
generating pseudo labels identifying the one or more anatomical objects in the one or more training medical images using the trained machine learning based model,
filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images, and
fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels.
2 . The computer-implemented method of claim 1 , wherein filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images comprises:
determining a feature representation of each of the generated pseudo labels using a feature extractor model of the machine learning based model;
determining a distance between a) the feature representation of each of the generated pseudo labels identifying the particular anatomical object in the one or more training medical images and b) the distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images; and
filtering the generated pseudo labels based on a comparison between the distances and a threshold.
3 . The computer-implemented method of claim 1 , wherein the machine learning based model is further trained by:
repeating the generating, the filtering, and the fine-tuning for a plurality of iterations.
4 . The computer-implemented method of claim 1 , wherein fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels comprises:
fine-tuning the trained machine learning based model to maximize an overlap between a predicted location of the one or more anatomical objects and a location defined in the ground truth labels.
5 . The computer-implemented method of claim 1 , wherein fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels comprises:
fine-tuning the trained machine learning based model to maximize an overlap between a predicted location of unlabeled anatomical objects and a location defined in the filtered pseudo labels for the unlabeled anatomical objects.
6 . The computer-implemented method of claim 1 , wherein the medical imaging analysis task is segmentation.
7 . An apparatus comprising:
means for receiving one or more input medical images;
means for performing a medical imaging analysis task based on the one or more input medical images using a machine learning based model; and
means for outputting results of the medical imaging analysis task,
wherein the machine learning based model is trained by:
receiving one or more training medical images and ground truth labels identifying one or more anatomical objects in the one or more training medical images,
training the machine learning based model for performing the medical imaging analysis task based on the one or more training medical images and the ground truth labels,
generating pseudo labels identifying the one or more anatomical objects in the one or more training medical images using the trained machine learning based model,
filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images, and
fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels.
8 . The apparatus of claim 7 , wherein filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images comprises:
determining a feature representation of each of the generated pseudo labels using a feature extractor model of the machine learning based model;
determining a distance between a) the feature representation of each of the generated pseudo labels identifying the particular anatomical object in the one or more training medical images and b) the distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images; and
filtering the generated pseudo labels based on a comparison between the distances and a threshold.
9 . The apparatus of claim 7 , wherein the machine learning based model is further trained by:
repeating the generating, the filtering, and the fine-tuning for a plurality of iterations.
10 . A non-transitory computer-readable storage medium comprising instructions which, when executed by a computer, cause the computer to carry out operations comprising:
receiving one or more input medical images;
performing a medical imaging analysis task based on the one or more input medical images using a machine learning based model; and
outputting results of the medical imaging analysis task,
wherein the machine learning based model is trained by:
receiving one or more training medical images and ground truth labels identifying one or more anatomical objects in the one or more training medical images,
training the machine learning based model for performing the medical imaging analysis task based on the one or more training medical images and the ground truth labels,
generating pseudo labels identifying the one or more anatomical objects in the one or more training medical images using the trained machine learning based model,
filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images, and
fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels.
11 . The non-transitory computer-readable storage medium of claim 10 , wherein fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels comprises:
fine-tuning the trained machine learning based model to maximize an overlap between a predicted location of the one or more anatomical objects and a location defined in the ground truth labels.
12 . The non-transitory computer-readable storage medium of claim 10 , wherein fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels comprises:
fine-tuning the trained machine learning based model to maximize an overlap between a predicted location of unlabeled anatomical objects and a location defined in the filtered pseudo labels for the unlabeled anatomical objects.
13 . The non-transitory computer-readable storage medium of claim 10 , wherein the medical imaging analysis task is segmentation.
14 . A computer-implemented method comprising:
receiving one or more training medical images and ground truth labels identifying one or more anatomical objects in the one or more training medical images;
training a machine learning based model for performing a medical imaging analysis task based on the one or more training medical images and the ground truth labels;
generating pseudo labels identifying the one or more anatomical objects in the one or more training medical images using the trained machine learning based model;
filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images;
fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels; and
outputting the fine-tuned machine learning based model.
15 . The computer-implemented method of claim 14 , wherein filtering the generated pseudo labels based on a comparison between 1) each of the generated pseudo labels identifying a particular anatomical object in the one or more training medical images and 2) a distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images comprises:
determining a feature representation of each of the generated pseudo labels using a feature extractor model of the machine learning based model;
determining a distance between a) the feature representation of each of the generated pseudo labels identifying the particular anatomical object in the one or more training medical images and b) the distribution of the ground truth labels identifying the particular anatomical object in the one or more training medical images; and
filtering the generated pseudo labels based on a comparison between the distances and a threshold.
16 . The computer-implemented method of claim 14 , wherein fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels comprises:
fine-tuning the trained machine learning based model to maximize an overlap between a predicted location of the one or more anatomical objects and a location defined in the ground truth labels.
17 . The computer-implemented method of claim 14 , wherein fine-tuning the trained machine learning based model for performing the medical imaging analysis task based on the one or more training medical images, the ground truth labels, and the filtered pseudo labels comprises:
fine-tuning the trained machine learning based model to maximize an overlap between a predicted location of unlabeled anatomical objects and a location defined in the filtered pseudo labels for the unlabeled anatomical objects.