methodw commited on
Commit
af11455
1 Parent(s): d98aa36

optimize index loading

Browse files
Files changed (1) hide show
  1. app.py +7 -7
app.py CHANGED
@@ -19,6 +19,13 @@ example_input = torch.rand(1, 3, 224, 224).to(device) # Adjust size if needed
19
  traced_model = torch.jit.trace(model, example_input)
20
  traced_model = traced_model.to(device)
21
 
 
 
 
 
 
 
 
22
 
23
  def process_image(image):
24
  """
@@ -28,10 +35,6 @@ def process_image(image):
28
  # This will include preprocessing the image, passing it through the model,
29
  # and then formatting the output (extracted features).
30
 
31
- # Load the index
32
- with open("xbgp-faiss-map.json", "r") as f:
33
- images = json.load(f)
34
-
35
  # Convert to RGB if it isn't already
36
  if image.mode != "RGB":
37
  image = image.convert("RGB")
@@ -62,7 +65,6 @@ def process_image(image):
62
  faiss.normalize_L2(vector)
63
 
64
  # Read the index file and perform search of top 50 images
65
- index = faiss.read_index("xbgp-faiss.index")
66
  distances, indices = index.search(vector, 50)
67
 
68
  matches = []
@@ -72,8 +74,6 @@ def process_image(image):
72
  gamerpic["id"] = images[matching_gamerpic]
73
  gamerpic["score"] = str(round((1 / (distances[0][idx] + 1) * 100), 2)) + "%"
74
 
75
- print(gamerpic)
76
-
77
  matches.append(gamerpic)
78
 
79
  return matches
 
19
  traced_model = torch.jit.trace(model, example_input)
20
  traced_model = traced_model.to(device)
21
 
22
+ # Load faiss index
23
+ index = faiss.read_index("xbgp-faiss.index")
24
+
25
+ # Load faiss map
26
+ with open("xbgp-faiss-map.json", "r") as f:
27
+ images = json.load(f)
28
+
29
 
30
  def process_image(image):
31
  """
 
35
  # This will include preprocessing the image, passing it through the model,
36
  # and then formatting the output (extracted features).
37
 
 
 
 
 
38
  # Convert to RGB if it isn't already
39
  if image.mode != "RGB":
40
  image = image.convert("RGB")
 
65
  faiss.normalize_L2(vector)
66
 
67
  # Read the index file and perform search of top 50 images
 
68
  distances, indices = index.search(vector, 50)
69
 
70
  matches = []
 
74
  gamerpic["id"] = images[matching_gamerpic]
75
  gamerpic["score"] = str(round((1 / (distances[0][idx] + 1) * 100), 2)) + "%"
76
 
 
 
77
  matches.append(gamerpic)
78
 
79
  return matches