System and method for semi-automated dataset annotation using similarity based clustering and in-context learning for segmentation
A method receiving a dataset including a plurality of images, executing an embedding model using the plurality of images to generate a plurality of image embeddings, clustering the plurality of image embeddings into a plurality of clusters based on similarities between the plurality of image embeddings, determining a density for each cluster based on an aggregate similarity of image embeddings of the cluster, based on the density for each of the plurality of clusters satisfying a criterion, for each cluster: receiving a first set of annotations for one or more pixels of each of a first subset of the images included in the cluster, inputting the annotated first subset of images into a vision transformer with a remaining unannotated subset of images not included in the annotated first subset, and generating a second plurality of annotations for the remaining unannotated subset of images not included in the annotated first subset.
1 . A method comprising:
receiving, by one or more processors, a dataset comprising a plurality of images;
executing, by the one or more processors, an embedding model using the plurality of images to generate a plurality of image embeddings, an image embedding for each of the plurality of images;
clustering, by the one or more processors, the plurality of image embeddings of the plurality of images into a plurality of clusters based on similarities between the plurality of image embeddings;
determining, by the one or more processors, a density for each of the plurality of clusters based on an aggregate similarity of image embeddings of the cluster; and
based on the density for each of the plurality of clusters satisfying a criterion, for each cluster of the plurality of clusters:
receiving, by the one or more processors, via a user input, a first set of annotations for one or more pixels of each of a first subset of the images included in the cluster;
inputting, by the one or more processors, the annotated first subset of images of the cluster into a vision transformer with a remaining unannotated subset of images of the cluster not included in the annotated first subset of images; and
generating, by the one or more processors, based on the input annotated first subset of images and the remaining unannotated subset of images input into the vision transformer, a second plurality of annotations for the remaining unannotated subset of images of the cluster not included in the annotated first subset of images.
2 . The method of claim 1 , wherein each of the image embeddings included in a particular cluster of the plurality of clusters each have a similarity score at or above a threshold value, the similarity score indicating a similarity of the image embedding to other image embeddings of the cluster.
3 . The method of claim 1 , further comprising:
comparing, by the one or more processors, each of the determined densities for each of the plurality of clusters to a predetermined threshold value; and
discarding, by the one or more processors, responsive to a determination that a density of a cluster is less than the predetermined threshold value, the cluster.
4 . The method of claim 3 , wherein discarding the cluster comprises:
preventing the images in the discarded cluster from being annotated by the user.
5 . The method of claim 1 , further comprising:
comparing, by the one or more processors, each of the determined densities for each of the plurality of clusters to a predetermined threshold value; and
automatically generating the second plurality of annotations for a particular cluster without user input, responsive to a determination that the density of the cluster is greater than or equal to the predetermined threshold value.
6 . The method of claim 1 , further comprising:
receiving, by the one or more processors, a user input regarding the second plurality of annotations;
updating, by the one or more processors, at least one of the second plurality of annotations based on the user input; and
updating, by the one or more processors, one or more images corresponding to the updated at least one of the second plurality of annotations based on the updated annotations.
7 . The method of claim 6 , wherein updating the at least one of the second plurality of annotations causes the vision transformer to reannotate the remaining images in the cluster.
8 . The method of claim 1 , wherein receiving the first set of annotations comprises:
receiving, via the user input, an indication from the user indicating a number of prompts to be used for a cluster;
displaying, to the user, a number of images included in the cluster corresponding to the number of prompts indicated by the user; and
receiving, via the user input, one or more annotations on each of the displayed images, the one or more annotations included in the first set of annotations.
9 . The method of claim 1 , wherein the density for each of the plurality of clusters satisfies the criterion based on a user input from a user.
10 . The method of claim 9 , further comprising, for each cluster of the plurality of clusters:
presenting the density to the user via a user interface;
receiving an indication of a number of clusters from the user interface via a user input; and
reclustering the plurality of image embeddings based on the received number of clusters.
11 . The method of claim 1 , further comprising:
responsive to generating the second plurality of annotations for a last cluster of the plurality of clusters, updating a database with the annotated plurality of images.
12 . A system comprising:
one or more non-transitory computer-readable media storing instructions thereon that, when executed by one or more processors, cause the one or more processors to:
receive a dataset comprising a plurality of images;
execute an embedding model using the plurality of images to generate a plurality of image embeddings, an image embedding for each of the plurality of images;
cluster the plurality of image embeddings of the plurality of images into a plurality of clusters based on similarities between the plurality of image embeddings;
determine a density for each of the plurality of clusters based on an aggregate similarity of image embeddings of the cluster; and
based on the density for each of the plurality of clusters satisfying a criterion, for each cluster of the plurality of clusters:
receive, via a user input, a first set of annotations for one or more pixels of each of a first subset of the images included in the cluster;
input the annotated first subset of images of the cluster into a vision transformer with a remaining unannotated subset of images of the cluster not included in the annotated first subset of images; and
generate, based on the input annotated first subset of images and the remaining unannotated subset of images input into the vision transformer, a second plurality of annotations for the remaining unannotated subset of images of the cluster not included in the annotated first subset of images.
13 . The system of claim 12 , wherein each of the image embeddings included in a particular cluster of the plurality of clusters each have a similarity score at or above a threshold value, the similarity score indicating a similarity of the image embedding to the remaining image embeddings.
14 . The system of claim 12 , wherein the instructions further cause the one or more processors to:
Compare each of the determined densities for each of the plurality of clusters to a predetermined threshold value.
15 . The system of claim 14 , wherein the instructions further cause the one or more processors to:
discard, responsive to a determination that a density of a cluster is less than the predetermined threshold value, the cluster.
16 . The system of claim 14 , wherein the second plurality of annotations for a particular cluster are generated automatically without user input responsive to a determination that the density of the cluster is greater than or equal to the predetermined threshold value.
17 . The system of claim 12 , wherein the instructions further cause the one or more processors to:
receive a user input regarding the second plurality of annotations;
update at least one of the second plurality of annotations based on the user input; and
update one or more images corresponding to the updated at least one of the second plurality of annotations based on the updated annotations.
18 . The system of claim 12 , wherein receiving the first set of annotations comprises:
receiving, via the user input, an indication from the user indicating a number of prompts to be used for a cluster;
displaying, to the user, a number of images included in the cluster corresponding to the number of prompts indicated by the user; and
receiving, via the user input, one or more annotations on each of the displayed images, the one or more annotations included in the first set of annotations.
19 . The system of claim 12 , wherein the density for each of the plurality of clusters satisfies the criterion based on a user input from a user, and wherein the instructions further cause the one or more processors to, for each cluster of the plurality of clusters:
present the density to the user via a user interface;
receive a number of clusters from the user interface via a user input; and
recluster the plurality of image embeddings based on the received number of clusters.
20 . One or more non-transitory computer-readable media storing instructions thereon that, when executed by one or more processors, cause the one or more processors to perform operations comprising:
receiving a dataset comprising a plurality of images;
executing an embedding model using the plurality of images to generate a plurality of image embeddings, an image embedding for each of the plurality of images;
clustering the plurality of image embeddings of the plurality of images into a plurality of clusters based on similarities between the plurality of image embeddings;
determining a density for each of the plurality of clusters based on an aggregate similarity of image embeddings of the cluster; and
based on the density for each of the plurality of clusters satisfying a criterion, for each cluster of the plurality of clusters:
receiving, via a user input, a first set of annotations for one or more pixels of each of a first subset of the images included in the cluster;
inputting the annotated first subset of images of the cluster into a vision transformer with a remaining unannotated subset of images of the cluster not included in the annotated first subset of images; and
generating, based on the input annotated first subset of images and the remaining unannotated subset of images input into the vision transformer, a second plurality of annotations for the remaining unannotated subset of images of the cluster not included in the annotated first subset of images.