Change pipeline tag, improve README, add LICENSE, make model faster and lighter

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Gabriele Berton, Carlo Masone
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,5 +1,5 @@
 ---
-pipeline_tag: visual-document-retrieval
+pipeline_tag: image-feature-extraction
 library_name: pytorch
 license: mit
 tags:
@@ -20,13 +20,24 @@ MegaLoc is an image retrieval model for visual place recognition (VPR) that achi
 
 ```python
 import torch
+import torchvision.transforms as tfm
+from PIL import Image
+
 model = torch.hub.load("gmberton/MegaLoc", "get_trained_model")
-model.eval()
 
-# Extract descriptor from an image
-image = torch.randn(1, 3, 322, 322)  # [B, 3, H, W] - any size works
-with torch.no_grad():
-    descriptor = model(image)  # [B, 8448] L2-normalized descriptor
+# Same preprocessing we use for evaluation: ImageNet normalization, resize to 322x322
+# (any resolution works, paper results are computed at 322x322)
+transform = tfm.Compose([
+    tfm.ToTensor(),
+    tfm.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    tfm.Resize(size=[322, 322], antialias=True),
+])
+
+images = torch.stack([transform(Image.open(path).convert("RGB")) for path in ["im1.jpg", "im2.jpg"]])
+with torch.inference_mode():
+    descriptors = model(images)  # shape [2, 8448], L2-normalized
+
+similarities = descriptors @ descriptors.T  # cosine similarities
 ```
 
 For benchmarking on VPR datasets, see [VPR-methods-evaluation](https://github.com/gmberton/VPR-methods-evaluation).

megaloc_model.py

@@ -142,13 +142,10 @@ class FeatureAggregator(nn.Module):
         p = torch.exp(p)
         p = p[:, :-1, :]
 
-        p = p.unsqueeze(1).repeat(1, self.cluster_dim, 1, 1)
-        f = f.unsqueeze(2).repeat(1, 1, self.num_clusters, 1)
-
         f = torch.cat(
             [
                 F.normalize(t, p=2, dim=-1),
-                F.normalize((f * p).sum(dim=-1), p=2, dim=1).flatten(1),
+                F.normalize(torch.einsum("bdn,bkn->bdk", f, p), p=2, dim=1).flatten(1),
             ],
             dim=-1,
         )
@@ -212,11 +209,8 @@ class MultiHeadAttention(nn.Module):
         qkv = qkv.permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
 
-        attn = (q @ k.transpose(-2, -1)) * self.scale
-        attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
-
-        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0)
+        x = x.transpose(1, 2).reshape(B, N, C)
         x = self.proj(x)
         x = self.proj_drop(x)