RRoundTable commited on
Commit
adb0128
·
1 Parent(s): 3d8379d

Use global variable for top k

Browse files
Files changed (1) hide show
  1. app.py +5 -4
app.py CHANGED
@@ -34,6 +34,8 @@ height, width = patch_size * patch_height_nums, patch_size * patch_width_nums
34
  # DINOV2
35
  model = torch.hub.load(DINOV2_REPO, DINOV2_MODEL)
36
 
 
 
37
 
38
  def read_image(image_path: str) -> np.ndarray:
39
  image = cv2.imread(image_path, cv2.IMREAD_COLOR)
@@ -104,8 +106,7 @@ def instance_recognition(
104
  image_embedding = infer(image)
105
 
106
  # search
107
- k = 5
108
- distances, indices = index.search(image_embedding, k)
109
 
110
  res = []
111
  for i in indices[0]:
@@ -134,10 +135,10 @@ with gr.Blocks() as demo:
134
 
135
  with gr.Row():
136
  output_images = [
137
- gr.Image(label=f"Similar {i + 1}") for i in range(5)
138
  ]
139
  distances = [
140
- gr.Text(label=f"Similar {i + 1} Distances") for i in range(5)
141
  ]
142
 
143
  infer_btn.click(
 
34
  # DINOV2
35
  model = torch.hub.load(DINOV2_REPO, DINOV2_MODEL)
36
 
37
+ # faiss
38
+ K= 5
39
 
40
  def read_image(image_path: str) -> np.ndarray:
41
  image = cv2.imread(image_path, cv2.IMREAD_COLOR)
 
106
  image_embedding = infer(image)
107
 
108
  # search
109
+ distances, indices = index.search(image_embedding, K)
 
110
 
111
  res = []
112
  for i in indices[0]:
 
135
 
136
  with gr.Row():
137
  output_images = [
138
+ gr.Image(label=f"Similar {i + 1}") for i in range(K)
139
  ]
140
  distances = [
141
+ gr.Text(label=f"Similar {i + 1} Distances") for i in range(K)
142
  ]
143
 
144
  infer_btn.click(