yahoojapan / NGT

Nearest Neighbor Search with Neighborhood Graph and Tree for High-dimensional Data
Apache License 2.0
1.24k stars 114 forks source link

Can I use it for Images? [How-To request] #85

Closed zubairahmed-ai closed 3 years ago

zubairahmed-ai commented 3 years ago

Hi Excellent work with your library, can you please show me how to use this with image similarity search with python examples?

Thanks

masajiro commented 3 years ago

Hi. The following is the code sample to build an index to search for images. Although this does not include searching, you can extract query image features in the same way as the following and search the index with them in the usual way.

import torch
import torchvision
import ngtpy
from torchvision import transforms as T
from PIL import Image
from pathlib import Path
image_dir = '/path/to/imagedir'
model = torchvision.models.resnet50(pretrained=True)
model = torch.nn.Sequential(*list(model.children())[:-1])
model.eval()
dim = 2048
transformer = T.Compose([
    T.Resize(256),
    T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
])
imagefiles = list(Path(image_dir).glob('*.jpg'))
ngtpy.create('anng', dim)
index = ngtpy.Index('anng')
with torch.no_grad():
    for imagefile in imagefiles:
        image = Image.open(imagefile).convert('RGB')
        x = transformer(image).unsqueeze(0)
        feat = model(x).view(-1, dim)
        feat = feat.numpy().tolist()[0]
        index.insert(feat)
index.build_index()
index.save()
zubairahmed-ai commented 3 years ago

Thank you, just in time when I am going to start working on this. I saw the other link you shared but this is good snippet above.