feat: output multiple images

app.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import torch
 import numpy as np
 import gradio as gr
@@ -5,46 +7,52 @@ from faiss import read_index
 from PIL import Image, ImageOps
 from datasets import load_dataset
 import torchvision.transforms as T
-from torchvision.models import resnet50
 
 from model import DINO
 
-transforms = T.Compose(
-    [T.ToTensor(), T.Resize(244), T.CenterCrop(224), T.Normalize([0.5], [0.5])]
-)
-
 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, transforms=transforms) -> torch.Tensor:
+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 = 1):
+def search_index(input_image, k: int = 1):
     with torch.no_grad():
-        embedding = model(augment(input_image))
+        embedding = model(augment(input_image).to(device))
         index = read_index("./bin/dino.index")
-        _, I = index.search(np.array(embedding[0].reshape(1, -1)), k)
-        indices = I[0]
-        answer = ""
-        for i, index in enumerate(indices[:1]):
+        _, 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"]
-    return retrieved_img 
+            images.append(retrieved_img)
+    return images
 
 
 app = gr.Interface(
     search_index,
-    inputs=gr.Image(),
-    outputs="image",
+    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__":

model.py

@@ -1,7 +1,9 @@
 import copy
 
+import torch
 from pytorch_lightning import LightningModule
 from torch import Tensor
+from torch.optim import SGD
 from torch.nn import Identity
 from torchvision.models import resnet50
 
@@ -13,11 +15,10 @@ from lightly.models.utils import (
     get_weight_decay_parameters,
     update_momentum,
 )
-from lightly.transforms import DINOTransform
 from lightly.utils.benchmarking import OnlineLinearClassifier
 from lightly.utils.scheduler import CosineWarmupScheduler, cosine_schedule
 
-from typing import  Union, Tuple, List
+from typing import Union, Tuple, List
 
 
 class DINO(LightningModule):