Single-Shot Learning for Object Re-Identification
I have put my initial idea to a test and used a small pretrained CNN (MobileNet) to compute features for reference images and stored the feature vectors in a &quot;database&quot;. Query images go through the exact same network and the resulting feature vector is used for nearest neighbor retrieval in the DB.
from glob import glob

import torch
from PIL import Image
from numpy.linalg import norm
from torchvision import transforms
from torchvision.models import mobilenet_v2

model = mobilenet_v2(pretrained=True)
model.eval()

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Generate DB
db = {}
dp_paths = glob('db/*.jpg')
for path in dp_paths:
    image = preprocess(Image.open(path)).unsqueeze(0)
    with torch.no_grad():
        output = model(image)
    db[output] = path

# Query
image = preprocess(Image.open('queries/box.jpg')).unsqueeze(0)
with torch.no_grad():
    query = model(image)

# Nearest Neighbor (poor man's version)
min_distance = float('inf')
candidate = None
for k, v in db.items():
    distance = norm(k.numpy() - query.numpy())
    if distance &lt; min_distance:
        min_distance = distance
        candidate = v

print(candidate, min_distance)

At least with my 5 test reference images and several query images it worked without a single failed &quot;classification&quot;. However, I am not sure if it will stand up to a larger test...