Clinic-driven multi-label classification framework for medical images
Machine-learning systems and methods can perform multi-label classification of medical images in a clinical context. A system can incorporate a triplet attention network that combines category-attention, self-attention, and cross-attention to learn high-quality label embeddings by mining effective information from medical images. The set of labels can include a single “hybrid” label assigned to multiple low-frequency diseases. The machine-learning system can be trained using a dual-pool contrastive learning technique. In the inference stage, the trained network can use a dual-pool contrastive inference technique to reduce the likelihood of false negatives and enhance the ability to detect diseases not seen in the training data.
1 . A computer-implemented method comprising:
defining a set of labels, wherein different labels in the set of labels correspond to different diseases;
obtaining a training data set, the training data set including a plurality of training images, wherein each training image in the training data set includes a medical image and a ground truth labeling that identifies the medical image as either positive or negative with respect to each label in the set of labels;
using the training data set to train a multi-label classifier network to assign labels from the set of labels to medical images, wherein the multi-label classifier network includes an image feature extractor, a label embedding extractor, a triplet attention transformer network, and a label prediction classifier, and wherein training the multi-label classifier network includes concurrently:
training the image feature extractor to extract image spatial features and category attention features from an input image;
training the label embedding extractor to generate initial label embeddings for the set of labels;
training the triple attention transformer network to generate updated label embeddings based on the image spatial features, the category attention features, and the initial label embeddings; and
training the label prediction classifier to predict a probability for each label in the set of labels based on the updated label embeddings, wherein a medical image is classified as positive or negative with respect to each label based on whether the probability exceeds a threshold; and
performing an inference process, wherein the inference process includes:
receiving a testing image; and
applying the trained multi-label classifier network to the testing image to produce an updated label embedding for the testing image and a probability of each label applying to the testing image.
2 . The method of claim 1 wherein defining the set of labels includes defining a single hybrid label that corresponds to a plurality of different diseases.
3 . The method of claim 2 wherein diseases are assigned to the hybrid label based on a low frequency of occurrence in the training data set.
4 . The method of claim 1 wherein the image feature extractor includes:
a convolutional backbone having a plurality of convolutional weights that outputs a feature matrix;
a reshaper that operates on a first matrix obtained by applying a shaping weight matrix to the feature matrix to produce the image spatial features; and
a category-wise attention module that operates on a second matrix obtained by applying an attention weight matrix to the feature matrix,
wherein training the multi-label classifier network includes training the convolutional weights, the shaping weight matrix, and the attention weight matrix.
5 . The method of claim 4 wherein the category-wise attention module includes a first neural network that implements a shortcut function and a second neural network to generate attention scores from the output of the first neural network.
6 . The method of claim 1 wherein the triplet attention transformer network includes:
an encoder module having at least two encoder self-attention layers, wherein the encoder module operates on an input comprising the image spatial features and a combination of the category attention features and the initial label embeddings;
a splitter module that splits an output of the encoder module into an updated image spatial features matrix and intermediate label embeddings;
a combiner module that combines the intermediate label embeddings and the category attention features; and
a decoder module having at least: a decoder self-attention layer that operates on an output of the combiner module; and at least one cross-attention layer that operates on the updated image spatial features matrix and an output of the decoder self-attention layer to produce the updated label embeddings.
7 . The method of claim 1 wherein training the multi-label classifier network includes minimizing a loss function, wherein the loss function includes:
a first term based on comparing a classification output of the multi-label classifier network with the ground truth labeling;
a second term based on comparing intermediate labels assigned based on the initial label embeddings with the ground truth labeling; and
a contrastive loss term.
8 . The method of claim 7 wherein the contrastive loss term includes:
an inter-pool contrastive loss term that measures a similarity of label embeddings for a given disease label between a positive sample pool that includes only training images for which the ground truth labeling is positive with respect to at least one of the labels and a negative sample pool that includes only training images for which the ground truth labeling is negative with respect to every one of the labels; and
an intra-pool contrastive loss term that measures a similarity of label embeddings for different labels within the positive sample pool.
9 . The method of claim 1 wherein performing the inference process further comprises:
applying a threshold to the probability of each label, wherein the testing image is identified as positive with respect to a given label if the probability of the given label exceeds the threshold and negative with respect to the given label otherwise.
10 . The method of claim 1 wherein performing the inference process further comprises:
producing an initial classifier output by applying a threshold to the probability of each label, wherein the initial classifier output identifies the testing image as positive with respect to a given label if the probability of the given label exceeds the threshold and negative with respect to the given label otherwise;
randomly selecting a negative sample pool that includes only training images for which the ground truth labeling is negative with respect to every one of the labels;
using the trained multi-label classifier network to generate a reference updated label embedding for each training image in the negative sample pool;
computing an average label embedding from the reference label embeddings;
computing a first similarity score for an updated label embedding generated by the trained multi-label classifier network for the testing image and the average label embedding;
computing a reference set of similarity scores that includes a similarity score for each of the reference label embeddings and the average label embedding; and
performing an outlier detection process on the first similarity score and the reference set of similarity scores to determine, for each label, whether a significant difference exists; and
generating a contrastive labeling for the testing image, wherein the contrastive labeling is positive with respect to each label where a significant difference exists and negative with respect to that label otherwise.
11 . The method of claim 10 wherein performing the inference process further comprises:
using the classification from the multi-label classifier network and the contrastive labeling to determine a final labeling for the testing image,
wherein the final labeling is negative with respect to a given label if both the classification from the multi-label classifier network and the contrastive labeling are negative and the final labeling is positive otherwise.
12 . A system comprising:
a memory; and
a processor coupled to the memory and configured to:
define a set of labels, wherein different labels in the set of labels correspond to different diseases;
obtain a training data set, the training data set including a plurality of training images, wherein each training image in the training data set includes a medical image and a ground truth labeling that identifies the medical image as either positive or negative with respect to each label in the set of labels;
use the training data set to train a multi-label classifier network to assign labels from the set of labels to medical images, wherein the multi-label classifier network includes an image feature extractor, a label embedding extractor, a triplet attention transformer network, and a label prediction classifier, and wherein training the multi-label classifier network includes concurrently:
training the image feature extractor to extract image spatial features and category attention features from an input image;
training the label embedding extractor to generate initial label embeddings for the set of labels;
training the triple attention transformer network to generate updated label embeddings based on the image spatial features, the category attention features, and the initial label embeddings; and
training the label prediction classifier to predict a probability for each label in the set of labels based on the updated label embeddings, wherein a medical image is classified as positive or negative with respect to each label based on whether the probability exceeds a threshold; and
perform an inference process, wherein the inference process includes:
receiving a testing image; and
applying the trained multi-label classifier network to the testing image to produce an updated label embedding for the testing image and a probability of each label applying to the testing image.
13 . The system of claim 12 wherein the image feature extractor includes:
a convolutional backbone having a plurality of convolutional weights that outputs a feature matrix;
a reshaper that operates on a first matrix obtained by applying a shaping weight matrix to the feature matrix to produce the image spatial features; and
a category-wise attention module that operates on a second matrix obtained by applying an attention weight matrix to the feature matrix,
wherein training the multi-label classifier network includes training the convolutional weights, the shaping weight matrix, and the attention weight matrix.
14 . The system of claim 13 wherein the category-wise attention module includes a first neural network that implements a shortcut function and a second neural network to generate attention scores from the output of the first neural network.
15 . The system of claim 12 wherein the triplet attention transformer network includes:
an encoder module having at least two encoder self-attention layers, wherein the encoder module operates on an input comprising the image spatial features and a combination of the category attention features and the initial label embeddings;
a splitter module that splits an output of the encoder module into an updated image spatial features matrix and intermediate label embeddings;
a combiner module that combines the intermediate label embeddings and the category attention features; and
a decoder module having at least: a decoder self-attention layer that operates on an output of the combiner module; and at least one cross-attention layer that operates on the updated image spatial features matrix and an output of the decoder self-attention layer to produce the updated label embeddings.
16 . A non-transitory computer-readable storage medium having stored therein program code instructions that, when executed by a processor in a computer system, cause the processor to perform a method comprising:
defining a set of labels, wherein different labels in the set of labels correspond to different diseases;
obtaining a training data set, the training data set including a plurality of training images, wherein each training image in the training data set includes a medical image and a ground truth labeling that identifies the medical image as either positive or negative with respect to each label in the set of labels;
using the training data set to train a multi-label classifier network to assign labels from the set of labels to medical images, wherein the multi-label classifier network includes an image feature extractor, a label embedding extractor, a triplet attention transformer network, and a label prediction classifier, and wherein training the multi-label classifier network includes concurrently:
training the image feature extractor to extract image spatial features and category attention features from an input image;
training the label embedding extractor to generate initial label embeddings for the set of labels;
training the triple attention transformer network to generate updated label embeddings based on the image spatial features, the category attention features, and the initial label embeddings; and
training the label prediction classifier to predict a probability for each label in the set of labels based on the updated label embeddings, wherein a medical image is classified as positive or negative with respect to each label based on whether the probability exceeds a threshold; and
performing an inference process, wherein the inference process includes:
receiving a testing image; and
applying the trained multi-label classifier network to the testing image to produce an updated label embedding for the testing image and a probability of each label applying to the testing image.
17 . The non-transitory computer-readable storage medium of claim 16 wherein defining the set of labels includes defining a single hybrid label that corresponds to a plurality of different diseases and wherein diseases are assigned to the hybrid label based on a low frequency of occurrence in the training data set.
18 . The non-transitory computer-readable storage medium of claim 16 wherein training the multi-label classifier network includes minimizing a loss function, wherein the loss function includes:
a first term based on comparing a classification output of the multi-label classifier network with the ground truth labeling;
a second term based on comparing intermediate labels assigned based on the initial label embeddings with the ground truth labeling; and
a contrastive loss term that includes:
an inter-pool contrastive loss term that measures a similarity of label embeddings for a given disease label between a positive sample pool that includes only training images for which the ground truth labeling is positive with respect to at least one of the labels and a negative sample pool that includes only training images for which the ground truth labeling is negative with respect to every one of the labels; and
an intra-pool contrastive loss term that measures a similarity of label embeddings for different labels within the positive sample pool.
19 . The non-transitory computer-readable storage medium of claim 16 wherein performing the inference process further comprises:
producing an initial classifier output by applying a threshold to the probability of each label, wherein the initial classifier output identifies the testing image as positive with respect to a given label if the probability of the given label exceeds the threshold and negative with respect to the given label otherwise;
randomly selecting a negative sample pool that includes only training images for which the ground truth labeling is negative with respect to every one of the labels;
using the trained multi-label classifier network to generate a reference updated label embedding for each training image in the negative sample pool;
computing an average label embedding from the reference label embeddings;
computing a first similarity score for an updated label embedding generated by the trained multi-label classifier network for the testing image and the average label embedding;
computing a reference set of similarity scores that includes a similarity score for each of the reference label embeddings and the average label embedding; and
performing an outlier detection process on the first similarity score and the reference set of similarity scores to determine, for each label, whether a significant difference exists; and
generating a contrastive labeling for the testing image, wherein the contrastive labeling is positive with respect to each label where a significant difference exists and negative with respect to that label otherwise.
20 . The non-transitory computer-readable storage medium of claim 19 wherein performing the inference process further comprises:
using the classification from the multi-label classifier network and the contrastive labeling to determine a final labeling for the testing image,
wherein the final labeling is negative with respect to a given label if both the classification from the multi-label classifier network and the contrastive labeling are negative and the final labeling is positive otherwise.