Merge pull request #44 from QVPR/issue_43_fix

Fixing val.py incorrect recalls issue #43

Author: Tobias-Fischer

patchnetvlad/training_tools/msls.py

@@ -99,6 +99,8 @@ def __init__(self, root_dir, cities='', nNeg=5, transform=None, mode='train', ta
         self.dbImages = []
         self.sideways = []
         self.night = []
+        self.qEndPosList = []
+        self.dbEndPosList = []
 
         self.all_pos_indices = []
 
@@ -186,6 +188,9 @@ def __init__(self, root_dir, cities='', nNeg=5, transform=None, mode='train', ta
                 self.qImages.extend(qSeqKeys)
                 self.dbImages.extend(dbSeqKeys)
 
+                self.qEndPosList.append(len(qSeqKeys))
+                self.dbEndPosList.append(len(dbSeqKeys))
+
                 qData = qData.loc[unique_qSeqIdx]
                 dbData = dbData.loc[unique_dbSeqIdx]
 

patchnetvlad/training_tools/val.py

@@ -82,11 +82,23 @@ def val(eval_set, model, encoder_dim, device, opt, config, writer, epoch_num=0,
     tqdm.write('====> Calculating recall @ N')
     n_values = [1, 5, 10, 20, 50, 100]
 
-    _, predictions = faiss_index.search(qFeat, max(n_values))
-
     # for each query get those within threshold distance
     gt = eval_set.all_pos_indices
 
+    # any combination of mapillary cities will work as a val set
+    qEndPosTot = 0
+    dbEndPosTot = 0
+    for cityNum, (qEndPos, dbEndPos) in enumerate(zip(eval_set.qEndPosList, eval_set.dbEndPosList)):
+        faiss_index = faiss.IndexFlatL2(pool_size)
+        faiss_index.add(dbFeat[dbEndPosTot:dbEndPosTot+dbEndPos, :])
+        _, preds = faiss_index.search(qFeat[qEndPosTot:qEndPosTot+qEndPos, :], max(n_values))
+        if cityNum == 0:
+            predictions = preds
+        else:
+            predictions = np.vstack((predictions, preds))
+        qEndPosTot += qEndPos
+        dbEndPosTot += dbEndPos
+
     correct_at_n = np.zeros(len(n_values))
     # TODO can we do this on the matrix in one go?
     for qIx, pred in enumerate(predictions):
@@ -104,4 +116,4 @@ def val(eval_set, model, encoder_dim, device, opt, config, writer, epoch_num=0,
         if write_tboard:
             writer.add_scalar('Val/Recall@' + str(n), recall_at_n[i], epoch_num)
 
-    return all_recalls
+    return all_recalls