Merge pull request #6 from prime-slam/salad-vpr

SALAD

aero_vloc/__init__.py

@@ -21,4 +21,4 @@
 from aero_vloc.primitives import Map, UAVSeq
 from aero_vloc.retrieval_system import RetrievalSystem
 from aero_vloc.utils import visualize_matches
-from aero_vloc.vpr_systems import AnyLoc, CosPlace, EigenPlaces, MixVPR, NetVLAD
+from aero_vloc.vpr_systems import AnyLoc, CosPlace, EigenPlaces, MixVPR, NetVLAD, SALAD

aero_vloc/retrieval_system.py

@@ -44,7 +44,6 @@ def __init__(
         for image in tqdm(
             sat_map, desc="Calculating of global descriptors for source DB"
         ):
-            image.descriptor = self.vpr_system.get_image_descriptor(image.path)
             self.global_descs.append(self.vpr_system.get_image_descriptor(image.path))
         self.index.create(np.asarray(self.global_descs))
 

aero_vloc/vpr_systems/__init__.py

@@ -16,4 +16,5 @@
 from aero_vloc.vpr_systems.eigenplaces import EigenPlaces
 from aero_vloc.vpr_systems.mixvpr import MixVPR
 from aero_vloc.vpr_systems.netvlad import NetVLAD
+from aero_vloc.vpr_systems.salad import SALAD
 from aero_vloc.vpr_systems.vpr_system import VPRSystem

aero_vloc/vpr_systems/salad/__init__.py

@@ -0,0 +1,14 @@
+#  Copyright (c) 2023, Ivan Moskalenko, Anastasiia Kornilova
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from aero_vloc.vpr_systems.salad.salad import SALAD

aero_vloc/vpr_systems/salad/salad.py

@@ -0,0 +1,52 @@
+#  Copyright (c) 2023, Ivan Moskalenko, Anastasiia Kornilova
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import torch
+
+from pathlib import Path
+from PIL import Image
+from torchvision import transforms as tvf
+
+from aero_vloc.utils import transform_image
+from aero_vloc.vpr_systems.vpr_system import VPRSystem
+
+
+class SALAD(VPRSystem):
+    """
+    Wrapper for [SALAD](https://github.com/serizba/salad) VPR method
+    """
+
+    def __init__(
+        self,
+        resize: int = 800,
+        gpu_index: int = 0,
+    ):
+        """
+        :param resize: The size to which the larger side of the image will be reduced while maintaining the aspect ratio
+        :param gpu_index: The index of the GPU to be used
+        """
+        super().__init__(gpu_index)
+        self.resize = resize
+        self.model = torch.hub.load("serizba/salad", "dinov2_salad")
+        self.model.eval().to(self.device)
+
+    def get_image_descriptor(self, image_path: Path):
+        image = Image.open(image_path).convert("RGB")
+        image = transform_image(image, self.resize).to(self.device)
+        _, h, w = image.shape
+        h_new, w_new = (h // 14) * 14, (w // 14) * 14
+        img_cropped = tvf.CenterCrop((h_new, w_new))(image)[None, ...]
+        with torch.no_grad():
+            descriptor = self.model(img_cropped)
+        descriptor = descriptor.cpu().numpy()[0]
+        return descriptor

requirements.txt

@@ -5,7 +5,9 @@ geopy==2.4.0
 numpy==1.26.1
 opencv_python==4.8.1.78
 Pillow==10.1.0
+prettytable==3.9.0
 pytorch_lightning==2.1.0
+pytorch-metric-learning==2.4.1
 Requests==2.31.0
 scikit_learn==1.3.2
 torch==2.1.0