Tensorflow
Image Similarity
Machine Learning
Computer Vision
Deep Learning

Tensorflow return similar images

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Returning similar images with TensorFlow usually means converting each image into a feature embedding and then comparing those embeddings with a similarity metric such as cosine similarity. The neural network is not directly answering "which image looks similar" in plain language. It is producing numeric representations that make similar images land near each other in feature space.

Build Embeddings With a Pretrained Model

A common approach is to use a pretrained CNN, remove its final classification layer, and treat the remaining output as an image embedding.

python
1import tensorflow as tf
2
3base_model = tf.keras.applications.MobileNetV2(
4    weights="imagenet",
5    include_top=False,
6    pooling="avg",
7)
8
9
10def load_image(path: str) -> tf.Tensor:
11    image = tf.io.read_file(path)
12    image = tf.image.decode_jpeg(image, channels=3)
13    image = tf.image.resize(image, (224, 224))
14    image = tf.keras.applications.mobilenet_v2.preprocess_input(image)
15    return image
16
17
18def embed_image(path: str) -> tf.Tensor:
19    image = load_image(path)
20    image = tf.expand_dims(image, axis=0)
21    embedding = base_model(image, training=False)
22    return tf.squeeze(embedding, axis=0)

This turns each image into a fixed-length vector that captures visual features useful for similarity search.

Compare With Cosine Similarity

Once you have embeddings, compare them numerically. Cosine similarity is common because it measures the angle between vectors rather than raw magnitude.

python
1import tensorflow as tf
2
3
4def cosine_similarity(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
5    a = tf.math.l2_normalize(a, axis=0)
6    b = tf.math.l2_normalize(b, axis=0)
7    return tf.reduce_sum(a * b)

A higher cosine similarity means the embeddings point in a more similar direction, which usually means the images are more visually alike according to the model.

Return the Top Similar Images

In a small image collection, you can compute embeddings once, store them, then compare a query image against all candidates.

python
1import tensorflow as tf
2
3
4def find_top_k_similar(query_path, candidate_paths, k=3):
5    query_embedding = embed_image(query_path)
6    scored = []
7
8    for path in candidate_paths:
9        candidate_embedding = embed_image(path)
10        score = float(cosine_similarity(query_embedding, candidate_embedding).numpy())
11        scored.append((path, score))
12
13    scored.sort(key=lambda item: item[1], reverse=True)
14    return scored[:k]
15
16
17results = find_top_k_similar(
18    "query.jpg",
19    ["cat1.jpg", "cat2.jpg", "dog1.jpg", "forest.jpg"],
20    k=2,
21)
22
23print(results)

This is enough for a prototype. For large datasets, precompute embeddings once and store them in a vector index instead of recomputing them for every query.

Why Classification Output Is Not Ideal

Developers sometimes try to compare the softmax outputs of an image classifier directly. That can work in narrow cases, but embeddings from an intermediate or pooled layer are usually better for similarity search because they capture richer structure than the final top-class probabilities.

In other words, similarity is usually an embedding problem, not a plain classification problem.

Practical Improvements

A few improvements make the system much more useful:

  • normalize all embeddings consistently
  • cache candidate embeddings instead of recalculating them
  • use approximate nearest neighbor search for large collections
  • fine-tune the embedding model on domain-specific images if needed

For example, fashion, medical, and product images often benefit from domain-specific training because generic ImageNet features may not capture the visual distinctions you care about most.

Common Pitfalls

A common mistake is comparing raw pixel arrays directly. That is extremely sensitive to lighting, cropping, and scale, and it usually performs poorly for semantic similarity.

Another issue is forgetting to preprocess images with the function expected by the chosen pretrained model. Wrong preprocessing means wrong embeddings.

Developers also sometimes recompute embeddings for the entire gallery on every query. That is fine for demos, but wasteful for real retrieval systems.

Finally, remember that "similar" depends on the embedding model. A model trained for general object recognition may rank images by object category, not by aesthetic style or fine-grained detail.

Summary

  • Use a pretrained CNN to convert images into feature embeddings.
  • Compare embeddings with cosine similarity or another vector metric.
  • Precompute candidate embeddings and rank them against the query embedding.
  • Prefer intermediate or pooled features over raw classifier outputs for similarity search.
  • For better results, adapt the embedding model to the visual domain you care about.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.