Closed zubairahmed-ai closed 4 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()
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.
Hi Excellent work with your library, can you please show me how to use this with image similarity search with python examples?
Thanks