app.py

#!/usr/bin/env python

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")

## Define Model and Dataset
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:
    img = Image.fromarray(img)
    if img.mode == "L":
        # Convert grayscale image to RGB by duplicating the single channel three times
        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, k: int = 1):
    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]
        images = []
        for i, 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(),
        gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Top K"),
    ],
    outputs=[
        gr.Gallery(label="Retrieved Images"),
    ],
)

if __name__ == "__main__":
    app.launch()