|
|
|
|
|
import torch |
|
import numpy as np |
|
import gradio as gr |
|
from faiss import read_index |
|
from PIL import Image, ImageOps |
|
from datasets import load_dataset |
|
import torchvision.transforms as T |
|
|
|
from model import DINO |
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
dataset = load_dataset("ethz/food101") |
|
model = DINO(batch_size_per_device=32, num_classes=1000).to(device) |
|
model.load_state_dict(torch.load("./bin/model.ckpt", map_location=device)["state_dict"]) |
|
|
|
|
|
def augment(img: np.ndarray) -> torch.Tensor: |
|
""" |
|
Helper Function to augment the image before we generate embeddings |
|
|
|
Args: |
|
img (np.ndarray): Input Image |
|
|
|
Returns: |
|
torch.Tensor |
|
""" |
|
img = Image.fromarray(img) |
|
if img.mode == "L": |
|
|
|
img = ImageOps.colorize(img, black="black", white="white") |
|
|
|
transforms = T.Compose( |
|
[T.ToTensor(), T.Resize(244), T.CenterCrop(224), T.Normalize([0.5], [0.5])] |
|
) |
|
|
|
return transforms(img).unsqueeze(0) |
|
|
|
|
|
def search_index(input_image: np.ndarray, k: int = 1) -> list: |
|
""" |
|
Retrieve the Top k images from the given input image |
|
|
|
Args: |
|
input_image (np.ndarray): Input Image |
|
k (int): number of images to fetch |
|
|
|
Returns: |
|
list: List of top k images retrieved using the embeddings |
|
generated from the input image |
|
""" |
|
images = [] |
|
|
|
with torch.no_grad(): |
|
embedding = model(augment(input_image).to(device)) |
|
index = read_index("./bin/dino.index") |
|
_, results = index.search(np.array(embedding[0].reshape(1, -1)), k) |
|
indices = results[0] |
|
for _, index in enumerate(indices[:k]): |
|
retrieved_img = dataset["train"][int(index)]["image"] |
|
images.append(retrieved_img) |
|
return images |
|
|
|
|
|
app = gr.Interface( |
|
search_index, |
|
inputs=[ |
|
gr.Image(label="Input Image"), |
|
gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Top K"), |
|
], |
|
outputs=[ |
|
gr.Gallery(label="Retrieved Images"), |
|
], |
|
) |
|
|
|
if __name__ == "__main__": |
|
app.launch() |
|
|