Automated image retrieval with graph neural network
A content retrieval system uses a graph neural network architecture to determine images relevant to an image designated in a query. The graph neural network learns a new descriptor space that can be used to map images in the repository to image descriptors and the query image to a query descriptor. The image descriptors characterize the images in the repository as vectors in the descriptor space, and the query descriptor characterizes the query image as a vector in the descriptor space. The content retrieval system obtains the query result by identifying a set of relevant images associated with image descriptors having above a similarity threshold with the query descriptor.
1. A computer system for automated image retrieval, the computer system comprising:
a processor configured to execute instructions; and
a non-transient computer-readable medium comprising instructions that when executed by the processor cause the processor to:
receive a request from a client device to retrieve images relevant to a query image;
access a set of trained weights for a set of neighbor nodes in an image retrieval graph of a query node associated with the query image, each weight of the set of trained weights representing an edge in the image retrieval graph connecting a respective neighbor node to the query node in the image retrieval graph;
generate a query descriptor mapping the query image to a descriptor space by applying the set of trained weights to combine outputs of the set of neighbor nodes at one or more layers of the image retrieval graph;
identify relevant images based on similarity of image descriptors associated with the relevant images to the query descriptor; and
return information about the relevant images to the client device.
2. The computer system of claim 1 , wherein instructions further cause the processor to:
sequentially generate outputs for at least a subset of nodes for the one or more layers of the image retrieval graph, wherein an output of a node in the subset of nodes at a current layer is generated by applying a set of weights for the current layer to outputs of neighbor nodes of the node in a previous layer.
3. The computer system of claim 2 , wherein the subset of nodes are first-order and second-order neighbor nodes of the query node.
4. The computer system of claim 1 , wherein the instructions further cause the processor to:
generate base descriptors for the set of nodes and a base descriptor for the query node;
generate adjacency scores between the query node and the set of nodes based on similarity between the base descriptor of the query node and the base descriptors of the set of nodes; and
identify the neighbor nodes of the query node as nodes that have adjacency scores above a threshold.
5. The computer system of claim 1 , wherein the query descriptor is generated by further weighting the neighbor nodes of the query node with corresponding adjacency scores for the neighbor nodes.
6. The computer system of claim 1 , wherein the neighbor nodes of the query node are first-order neighbors of the query node.
7. The computer system of claim 1 , wherein the query descriptor is an output of the query node at a last layer of the image retrieval graph.
8. The computer system of claim 1 , wherein the set of weights is trained by the process of:
initializing an estimated set of weights; and
repeatedly performing the steps of:
generating estimated image descriptors for the set of nodes by applying the estimated set of weights to the set of nodes at the one or more layers of the image retrieval graph;
determining a loss function as a combination of losses for the set of nodes, a loss for a node indicating a similarity between an estimated image descriptor for the node and estimated image descriptors for neighbor nodes of the node; and
updating the estimated set of weights to reduce the loss function.
9. A method for automated image retrieval, comprising:
receiving a request from a client device to retrieve images relevant to a query image;
accessing a set of trained weights for a set of neighbor nodes in an image retrieval graph of a query node associated with the query image, each weight of the set of trained weights representing an edge in the image retrieval graph connecting a respective neighbor node to the query node in the image retrieval graph;
generating a query descriptor mapping the query image to a descriptor space by applying the set of trained weights to combine outputs of the set of neighbor nodes at the one or more layers of the image retrieval graph;
identifying relevant images based on similarity of image descriptors associated with the relevant images to the query descriptor; and
returning information about the relevant images to the client device.
10. The method of claim 9 , further comprising:
sequentially generating outputs for at least a subset of nodes for the one or more layers of the image retrieval graph, wherein an output of a node in the subset of nodes at a current layer is generated by applying a set of weights for the current layer to outputs of neighbor nodes of the node in a previous layer.
11. The method of claim 10 , wherein the subset of nodes are first-order and second-order neighbor nodes of the query node.
12. The method of claim 9 , the method further comprising:
generating base descriptors for the set of nodes and a base descriptor for the query node;
generating adjacency scores between the query node and the set of nodes based on similarity between the base descriptor of the query node and the base descriptors of the set of nodes; and
identifying the neighbor nodes of the query node as nodes that have adjacency scores above a threshold.
13. The method of claim 9 , wherein the query descriptor is generated by further weighting the neighbor nodes of the query node with corresponding adjacency scores for the neighbor nodes.
14. The method of claim 9 , wherein the neighbor nodes of the query node are first-order neighbors of the query node.
15. The method of claim 9 , wherein the query descriptor is an output of the query node at a last layer of the image retrieval graph.
16. The method of claim 9 , wherein the set of weights in the image retrieval graph is trained by the process of:
initializing an estimated set of weights; and
repeatedly performing the steps of:
generating estimated image descriptors for the set of nodes by applying the estimated set of weights to the set of nodes at the one or more layers of the image retrieval graph;
determining a loss function as a combination of losses for the set of nodes, a loss for a node indicating a similarity between an estimated image descriptor for the node and estimated image descriptors for neighbor nodes of the node; and
updating the estimated set of weights to reduce the loss function.