Spaces:
Runtime error
Runtime error
import torch | |
from torch import optim, nn | |
from torchvision import models, transforms | |
from torchvision.models.vgg import VGG16_Weights | |
class FeatureExtractor(nn.Module): | |
def __init__(self, model): | |
super(FeatureExtractor, self).__init__() | |
# Extract VGG-16 Feature Layers | |
self.features = list(model.features) | |
self.features = nn.Sequential(*self.features) | |
# Extract VGG-16 Average Pooling Layer | |
self.pooling = model.avgpool | |
# Convert the image into one-dimensional vector | |
self.flatten = nn.Flatten() | |
# Extract the first part of fully-connected layer from VGG16 | |
self.fc = model.classifier[0] | |
def forward(self, x): | |
# It will take the input 'x' until it returns the feature vector called 'out' | |
out = self.features(x) | |
out = self.pooling(out) | |
out = self.flatten(out) | |
out = self.fc(out) | |
return out | |
# Initialize the model | |
model = models.vgg16(weights=VGG16_Weights.DEFAULT) | |
new_model = FeatureExtractor(model) | |
# Change the device to GPU | |
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu") | |
new_model = new_model.to(device) | |
IMG_RESIZE_SIZE = 224 | |
IMG_PATH = "data" | |
import cv2 | |
from tqdm import tqdm | |
import numpy as np | |
# Transform the image, so it becomes readable with the model | |
transform = transforms.Compose([ | |
transforms.ToPILImage(), | |
transforms.Resize((IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)), | |
transforms.ToTensor() | |
]) | |
# Will contain the feature | |
features = [] | |
mappings = {} | |
import glob | |
# files = glob.glob("/content/drive/Shareddrives/ndl/kao/dataset 3/*.jpg") | |
files = glob.glob(f"{IMG_PATH}/*.jpg") | |
files.sort() | |
for index in tqdm(range(len(files))): | |
path = files[index] | |
img = cv2.imread(path) | |
# Transform the image | |
img = transform(img) | |
# Reshape the image. PyTorch model reads 4-dimensional tensor | |
# [batch_size, channels, width, height] | |
# img = img.reshape(1, 3, 448, 448) | |
img = img.reshape(1, 3, IMG_RESIZE_SIZE, IMG_RESIZE_SIZE) | |
img = img.to(device) | |
# We only extract features, so we don't need gradient | |
with torch.no_grad(): | |
# Extract the feature from the image | |
feature = new_model(img) | |
# Convert to NumPy Array, Reshape it, and save it to features variable | |
features.append(feature.cpu().detach().numpy().reshape(-1)) | |
mappings[index] = { | |
"nconst": path.split("/")[-1].split(".")[0], | |
"name": "", | |
"url": "" | |
} | |
# Convert to NumPy Array | |
features = np.array(features) | |
import json | |
with open('mappings.json', mode='wt', encoding='utf-8') as file: | |
json.dump(mappings, file, ensure_ascii=False, indent=2) | |
N_TREES = 1000 | |
from annoy import AnnoyIndex | |
annoy_index = AnnoyIndex(features.shape[1], metric='euclidean') | |
for i in range(len(features)): | |
feature = features[i] | |
annoy_index.add_item(i, feature) | |
# k-d tree γγγ«γγγ | |
annoy_index.build(n_trees=N_TREES) | |
annoy_index.save("../models/index.ann") |