Pawel Piwowarski commited on
Commit
09a106d
·
1 Parent(s): 46c6bf9

added query rotation

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -449,6 +449,9 @@ def search_and_retrieve(
449
  tfm.Resize((MATCHING_IMG_SIZE, MATCHING_IMG_SIZE), antialias=True)(query_image)
450
  )
451
 
 
 
 
452
  progress(0.4, desc="Processing candidates")
453
  for faiss_idx in flat_indices:
454
 
@@ -456,6 +459,7 @@ def search_and_retrieve(
456
 
457
 
458
  best_rotation_index = faiss_idx // num_db_images
 
459
 
460
  if image_index in processed_image_indices:
461
  continue
@@ -513,9 +517,7 @@ def search_and_retrieve(
513
  candidate_img_tensor = tfm.Resize(
514
  (MATCHING_IMG_SIZE * 3, MATCHING_IMG_SIZE * 3), antialias=True
515
  )(candidate_img_tensor)
516
- candidate_img_tensor = F.rotate(
517
- candidate_img_tensor, [0, 90, 180, 270][best_rotation_index]
518
- )
519
 
520
  candidate_img_tensor = candidate_img_tensor.to(DEVICE)
521
 
 
449
  tfm.Resize((MATCHING_IMG_SIZE, MATCHING_IMG_SIZE), antialias=True)(query_image)
450
  )
451
 
452
+
453
+
454
+
455
  progress(0.4, desc="Processing candidates")
456
  for faiss_idx in flat_indices:
457
 
 
459
 
460
 
461
  best_rotation_index = faiss_idx // num_db_images
462
+ query_tensor = F.rotate(query_tensor, [0, 90, 180, 270][best_rotation_index] ) * -1
463
 
464
  if image_index in processed_image_indices:
465
  continue
 
517
  candidate_img_tensor = tfm.Resize(
518
  (MATCHING_IMG_SIZE * 3, MATCHING_IMG_SIZE * 3), antialias=True
519
  )(candidate_img_tensor)
520
+
 
 
521
 
522
  candidate_img_tensor = candidate_img_tensor.to(DEVICE)
523