Image Search
Finding similar photos instantly.
An expert spots a strange new mark on an apple, and asks, have we ever seen this before? You've got millions of unlabelled photos. Text search can't help, because nobody tagged them. But a trained network has already learned to turn images into meaningful numbers.
- A strange new mark
- Have we seen this before?
- No tags, just numbers
Pretrained embeddings
We take a network trained on ImageNet, and remove its final layer. What's left turns an image into a list of five hundred and twelve numbers. Similar images get similar numbers. We embed three hundred patches, then for each one, we find its nearest neighbour. Is that neighbour the same type, bruised or healthy? Run it. Over ninety nine percent. Remember, this network has never seen an apple bruise. It learned its features from everyday objects, and they still work here. Store those vectors in a vector database, and you can search millions of images by example.
import torch
from torchvision import models
from patch import make
m = models.resnet18(weights="DEFAULT").eval()
m.fc = torch.nn.Identity()
X, y = make(300, bruise_rate=.5, seed=2)
with torch.no_grad():
e = torch.nn.functional.normalize(m(torch.nn.functional.interpolate(X, 96)), dim=1)
sim = e @ e.T
sim.fill_diagonal_(-1)
print("Neighbour has the same label:", round((y[sim.argmax(1)] == y).float().mean().item(), 3))Neighbour has the same label: 0.997
Also uses belt.py
import numpy as np
import cv2
def belt(n=7, seed=1, size=(480, 640), bruises=2, light=1.0):
"""Draw a simulated conveyor belt of apples. Returns the image and the true boxes."""
rng = np.random.default_rng(seed)
h, w = size
img = np.full((h, w, 3), (70, 75, 80), np.uint8)
img = cv2.add(img, rng.normal(0, 6, img.shape).clip(0, 30).astype(np.uint8))
boxes, spots, taken = [], [], []
while len(boxes) < n:
r = int(rng.integers(32, 48))
x, y = int(rng.integers(r + 5, w - r - 5)), int(rng.integers(r + 5, h - r - 5))
if any((x - a) ** 2 + (y - b) ** 2 < (r + c + 8) ** 2 for a, b, c in taken):
continue
taken.append((x, y, r))
red = int(rng.integers(150, 225))
cv2.circle(img, (x, y), r, (35, 40, red), -1, cv2.LINE_AA)
cv2.circle(img, (x - r // 3, y - r // 3), r // 4, (90, 90, min(255, red + 30)), -1, cv2.LINE_AA)
boxes.append((x - r, y - r, x + r, y + r))
for x, y, r in taken[:bruises]:
cv2.ellipse(img, (x + r // 3, y + r // 4), (r // 3, r // 4), 20, 0, 360, (25, 45, 70), -1, cv2.LINE_AA)
spots.append((x, y, r))
img = cv2.GaussianBlur(img, (3, 3), 0)
return np.clip(img * light, 0, 255).astype(np.uint8), boxes, spotsAlso uses patch.py
import numpy as np
import cv2
import torch
def make(n, seed=0, light=1.0, bruise_rate=0.5):
"""Simulated apple photos. Label 1 means bruised."""
rng = np.random.default_rng(seed)
X, y = [], []
for _ in range(n):
img = np.full((48, 48, 3), (70, 75, 80), np.uint8)
c = tuple(int(v) for v in rng.integers(19, 30, 2))
r = int(rng.integers(15, 20))
cv2.circle(img, c, r, (35, 40, int(rng.integers(150, 225))), -1, cv2.LINE_AA)
bad = rng.random() < bruise_rate
if bad:
at = (c[0] + int(rng.integers(-8, 9)), c[1] + int(rng.integers(-8, 9)))
cv2.circle(img, at, int(rng.integers(3, 6)), (25, 45, 70), -1, cv2.LINE_AA)
img = img * light + rng.normal(0, 8, img.shape)
X.append(img.clip(0, 255) / 255)
y.append(int(bad))
X = torch.tensor(np.array(X), dtype=torch.float32).permute(0, 3, 1, 2)
return X, torch.tensor(y)
def net():
return torch.nn.Sequential(
torch.nn.Conv2d(3, 16, 3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2),
torch.nn.Conv2d(16, 32, 3, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(2),
torch.nn.Flatten(), torch.nn.Linear(32 * 12 * 12, 2))
def train(model, X, y, epochs=8):
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
for _ in range(epochs):
for i in torch.randperm(len(X)).split(64):
loss = torch.nn.functional.cross_entropy(model(X[i]), y[i])
opt.zero_grad(); loss.backward(); opt.step()
return model
def accuracy(model, X, y):
return (model(X).argmax(1) == y).float().mean().item()What you can do with it
Embeddings unlock a lot. Finding similar products in shops. Removing near duplicates. Picking a varied set of images to label. Spotting drift, by watching where new images land. And it's quick, because you only embed each image once. With models such as CLIP, you can even search with words. Show me bruised red apples, and it finds them, with no labels.
Summary
Networks turn images into searchable numbers. Next, a test of how far general models have come.
Key takeaways- Remove the last layer
- Vectors capture appearance
- Works on new subjects