Update main.py

main.py

@@ -13,40 +13,40 @@ from PIL import Image
 import matplotlib.cm as cm
 import matplotlib as mpl
 
-# img_names = []
-# with open('list_eval_partition.txt', 'r') as f:
-#     for line in tqdm.tqdm(f):
-#         img_name, dtype = line.rstrip().split(' ')
-#         img_names.append(img_name)
+img_names = []
+with open('list_eval_partition.txt', 'r') as f:
+    for line in f:
+        img_name, dtype = line.rstrip().split(' ')
+        img_names.append(img_name)
 
 
 # For a model pretrained on VGGFace2
 print('Loading model weights ........')
 
-# class SiameseModel(nn.Module):
-#     def __init__(self):
-#         super().__init__()
-#         self.backbone = InceptionResnetV1(pretrained='vggface2')
-#     def forward(self, x):
-#         x = self.backbone(x)
-#         x = torch.nn.functional.normalize(x, dim=1)
-#         return x
+class SiameseModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = InceptionResnetV1(pretrained='vggface2')
+    def forward(self, x):
+        x = self.backbone(x)
+        x = torch.nn.functional.normalize(x, dim=1)
+        return x
     
-# model = SiameseModel()
-# model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
-# model.eval()
+model = SiameseModel()
+model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
+model.eval()
 
 
 # Make FAISS index
-# print('Make index .............')
-# index = faiss.IndexFlatL2(512)
+print('Make index .............')
+index = faiss.IndexFlatL2(512)
 
-# hf = h5py.File('face_vecs_full.h5', 'r')
-# for key in tqdm.tqdm(hf.keys()):
-#     vec = np.array(hf.get(key))
-#     index.add(vec)
+hf = h5py.File('face_vecs_full.h5', 'r')
+for key in tqdm.tqdm(hf.keys()):
+    vec = np.array(hf.get(key))
+    index.add(vec)
 
-# hf.close()
+hf.close()
 
 print("Finished indexing")
 
@@ -98,16 +98,16 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
         
         gallery = gr.Gallery(label='Retrieval Images', columns=[5], show_label=True, scale=2)
         
-    # img_dir = './img_align_celeba'    
-    # examples = random.choices(img_names, k=6)
-    # examples = [os.path.join(img_dir, ex) for ex in examples]
-    # examples = [Image.open(img) for img in examples]
+    img_dir = './img_align_celeba'    
+    examples = random.choices(img_names, k=6)
+    examples = [os.path.join(img_dir, ex) for ex in examples]
+    examples = [Image.open(img) for img in examples]
                                   
-    # with gr.Row():
-    #     gr.Examples(
-    #         examples = examples,
-    #         inputs = image
-    #     )
+    with gr.Row():
+        gr.Examples(
+            examples = examples,
+            inputs = image
+        )
             
             
     btn.click(image_search,