Spaces:
Runtime error
Runtime error
File size: 3,049 Bytes
cab1b96 cca6c8d cab1b96 dd3ffb7 cab1b96 f726e05 cab1b96 f726e05 cab1b96 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import annoy
from typing import Tuple
import json
EMBEDDING_DIMENSION = 4096
ANNOY_INDEX_FILE = "models/index.ann"
ANNOY_MAPPING_FILE = "models/mappings.json"
IMG_RESIZE_SIZE = 224
####
def load_annoy_index(
index_file=ANNOY_INDEX_FILE,
mapping_file=ANNOY_MAPPING_FILE,
) -> Tuple[annoy.AnnoyIndex, dict]:
"""Load annoy index and associated mapping file"""
annoy_index = annoy.AnnoyIndex(f=EMBEDDING_DIMENSION, metric='angular')
annoy_index.load(index_file)
with open(ANNOY_MAPPING_FILE) as f:
mappings = json.load(f)
with open(mapping_file) as f:
mapping = json.load(f)
mapping = {int(k): v for k, v in mapping.items()}
return annoy_index, mappings
###
import torch
from torch import nn
from torchvision import models, transforms
# from torchvision.models.vgg import VGG16_Weights
# Transform the image, so it becomes readable with the model
transform = transforms.Compose([
transforms.ToPILImage(),
# transforms.CenterCrop(512),
# transforms.Resize(448),
transforms.Resize((IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)),
transforms.ToTensor()
])
import cv2
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)
model = models.vgg16(pretrained=True)
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)
import PIL
def analyze_image(
image, annoy_index, n_matches: int = 1, num_jitters: int = 1, model: str = "large"
):
PIL.Image.fromarray(image).save("input_img.png")
img = cv2.imread("input_img.png")
# 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
v = feature.cpu().detach().numpy().reshape(-1)
results = annoy_index.get_nns_by_vector(v, n_matches, include_distances=True)
return results |