haihuynh commited on
Commit
0272604
1 Parent(s): 4e0f2a8

Update cbir_system.py

Browse files
Files changed (1) hide show
  1. cbir_system.py +13 -11
cbir_system.py CHANGED
@@ -22,18 +22,18 @@ with open('list_eval_partition.txt', 'r') as f:
22
  # For a model pretrained on VGGFace2
23
  print('Loading model weights ........')
24
 
25
- class SiameseModel(nn.Module):
26
- def __init__(self):
27
- super().__init__()
28
- self.backbone = InceptionResnetV1(pretrained='vggface2')
29
- def forward(self, x):
30
- x = self.backbone(x)
31
- x = torch.nn.functional.normalize(x, dim=1)
32
- return x
33
 
34
- model = SiameseModel()
35
- model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
36
- model.eval()
37
 
38
 
39
  # Make FAISS index
@@ -47,6 +47,8 @@ for key in hf.keys():
47
 
48
  hf.close()
49
 
 
 
50
  # Function to search image
51
  def image_search(image, k=5):
52
 
 
22
  # For a model pretrained on VGGFace2
23
  print('Loading model weights ........')
24
 
25
+ # class SiameseModel(nn.Module):
26
+ # def __init__(self):
27
+ # super().__init__()
28
+ # self.backbone = InceptionResnetV1(pretrained='vggface2')
29
+ # def forward(self, x):
30
+ # x = self.backbone(x)
31
+ # x = torch.nn.functional.normalize(x, dim=1)
32
+ # return x
33
 
34
+ # model = SiameseModel()
35
+ # model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
36
+ # model.eval()
37
 
38
 
39
  # Make FAISS index
 
47
 
48
  hf.close()
49
 
50
+ print("Finished indexing")
51
+
52
  # Function to search image
53
  def image_search(image, k=5):
54