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;
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 an image retrieval graph including one or more layers with a set of trained weights, the image retrieval graph having a set of nodes representing a set of images and a set of edges representing connections between the set of images, wherein a node is associated with an image descriptor mapping the corresponding image represented by the node to a descriptor space;
generate a query node representing the query image in the image retrieval graph;
generate a second set of edges representing connections between the query node and neighbor nodes of the query node;
generate a query descriptor mapping the query image to the descriptor space by applying the set of weights to outputs of the neighbor nodes of the query node at the one or more layers of the image retrieval graph;
identify the images relevant to the query image by selecting a relevant subset of nodes, the image descriptors for the relevant subset of nodes having above a similarity threshold with the query descriptor; and
return the images represented by the relevant subset of nodes as a query result to the client device.
2. The computer system of claim 1 , wherein instructions to generate the query descriptor further cause the processor to:
sequentially generate outputs for at least a second subset of nodes for the one or more layers of the image retrieval graph, wherein an output of a node in the second 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 second 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 obtained as the 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 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.
9. A method for automated image retrieval, comprising:
receiving a request from a client device to retrieve images relevant to a query image;
accessing an image retrieval graph including one or more layers, the image retrieval graph having a set of image nodes representing a set of images and a set of edges representing connections between the set of images, wherein the one or more layers are associated with a set of weights, and wherein an image node is associated with an image descriptor mapping the corresponding image represented by the image node to a descriptor space;
generating a query node representing the query image in the image retrieval graph;
generating a second set of edges representing connections between the query node and at least a subset of image nodes that are identified as neighbor nodes of the query image;
generating a query descriptor mapping the query image to the descriptor space by applying the set of weights to outputs of the neighbor nodes and the query node at the one or more layers of the image retrieval graph;
identifying the images relevant to the query image by selecting a second subset of image nodes, the image descriptors for the second subset of image nodes having above a similarity threshold with the query descriptor; and
returning the images represented by the second subset of image nodes as a query result to the client device.
10. The method of claim 9 , wherein generating the query descriptor further comprises:
sequentially generating outputs for at least a second subset of nodes for the one or more layers of the image retrieval graph, wherein an output of a node in the second 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 second 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 obtained as the 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.