Update benchmark with full 6-model comparison

megaloc/README.md

@@ -1,80 +1,121 @@
-# MegaLoc Quantized Models
+# MegaLoc
+An image retrieval model for any localization task, which achieves SOTA on most VPR datasets, including indoor and outdoor ones.
 
-Quantized ONNX and CoreML models for [MegaLoc](https://arxiv.org/abs/2502.17237) — SOTA visual place recognition.
+[Gradio Demo](https://gmberton.github.io/demos-url/megaloc/) - [ArXiv](https://arxiv.org/abs/2502.17237) - [Paper on ArXiv](https://arxiv.org/pdf/2502.17237) - [Paper on HF](https://huggingface.co/papers/2502.17237) - [Model on HF](https://huggingface.co/gberton/MegaLoc).
 
-## Directory
+### Demo
+Try the demo on your own images to see how good MegaLoc is! The demo uses a database of ~5M street-view images from San Francisco, and when you upload one it will find the most similar one from the same place.
 
+<img width="746" height="576" alt="image" src="https://github.com/user-attachments/assets/4e7a3eec-dfee-4aae-83cc-f5146a1b421d" />
+
+
+### Using the model
+You can use the model with torch.hub, as simple as this
 ```
-megaloc/
-├── onnx/         # ONNX models (cross-platform)
-│   ├── megaloc.onnx           # FP32 (873 MB, batch=1)
-│   ├── megaloc_fp16.onnx      # FP16 (437 MB, dynamic batch)
-│   └── megaloc_int8.onnx      # INT8 (223 MB, batch=1)
-└── coreml/       # CoreML models (macOS / Apple Silicon, all multi-batch)
-    ├── megaloc_coreml_fp32.mlpackage   # FP32 (873 MB)
-    ├── megaloc_coreml_fp16.mlpackage   # FP16 (437 MB)
-    └── megaloc_coreml_int8.mlpackage   # INT8 (219 MB)
+import torch
+model = torch.hub.load("gmberton/MegaLoc", "get_trained_model")
 ```
 
-## Benchmark (MacBook Pro M-series, 518×518 input)
+For more complex uses, like computing results on VPR datasets, visualizing predictions and so on, you can use our [VPR-methods-evaluation](https://github.com/gmberton/VPR-methods-evaluation), which lets you do all this for MegaLoc and multiple other VPR methods on labelled or unlabelled datasets.
 
-### ONNX (CPU)
+### Qualitative examples
+Here are some examples of top-1 retrieved images from the SF-XL test set, which has 2.8M images as database.
 
-| Model | Size | CosSim | B=1 | B=2 | B=4 | B=8 |
-|-------|------|--------|-----|-----|-----|-----|
-| FP32 | 873 MB | 1.000000 | 399ms / 2.5fps | 872ms / 2.3fps | 1514ms / 2.6fps | 2699ms / 3.0fps |
-| FP16 | 437 MB | 0.999989 | 400ms / 2.5fps | 850ms / 2.4fps | 1602ms / 2.5fps | 2876ms / 2.8fps |
-| INT8 | 223 MB | 0.996453 | 475ms / 2.1fps | — | — | — |
+![teaser](https://github.com/user-attachments/assets/a90b8d4c-ab53-4151-aacc-93493d583713)
 
-FP16 is near-lossless at all batch sizes. INT8 preserves Recall@1 = 100% in retrieval tasks.
 
-### CoreML (Apple Silicon GPU, multi-batch)
 
-| Model | Size | CosSim | B=1 latency | Notes |
-|-------|------|--------|-------------|-------|
-| FP32 | 873 MB | 1.000000 | ~74ms / 13fps | RangeDim(1,16) dynamic batch |
-| FP16 | 437 MB | TBD | TBD | RangeDim(1,16) dynamic batch |
-| INT8 | 219 MB | TBD | TBD | RangeDim(1,16) dynamic batch |
+---
 
-CoreML is ~5× faster than ONNX Runtime (GPU vs CPU). All CoreML models support dynamic batch via `RangeDim(1, 16)`.
+## Available Models
 
-## Usage
+Pre-quantized models on HuggingFace: [Realcat/image_retrieval_checkpoints](https://huggingface.co/Realcat/image_retrieval_checkpoints)
 
-### ONNX
+```bash
+# ONNX
+huggingface-cli download Realcat/image_retrieval_checkpoints --include "megaloc/onnx/*" --local-dir .
+# CoreML
+huggingface-cli download Realcat/image_retrieval_checkpoints --include "megaloc/coreml/*" --local-dir .
+```
 
-```python
-import onnxruntime as ort, numpy as np
-from PIL import Image
+### Benchmark (MacBook Pro M-series, 518×518, B=1 vs PyTorch FP32)
 
-sess = ort.InferenceSession('megaloc_fp16.onnx', providers=['CPUExecutionProvider'])
+| Model | Size | CosSim | SNR | Latency | FPS | Multi-Batch |
+|-------|------|--------|-----|---------|-----|-------------|
+| **PyTorch FP32** | (mem) | 1.000000 | ∞ | 367 ms | 2.7 | — |
+| ONNX FP32 | 873 MB | 1.000000 | 80.5 dB | 367 ms | 2.7 | **Y** |
+| ONNX FP16 | 437 MB | 0.999989 | 46.4 dB | 345 ms | 2.9 | **Y** |
+| ONNX INT8 | 223 MB | 0.996625 | 21.7 dB | 395 ms | 2.5 | **Y** |
+| CoreML FP32 | 873 MB | 1.000000 | 80.5 dB | **74 ms** | **13.6** | N |
+| CoreML INT8 | 219 MB | 0.998699 | 25.8 dB | **74 ms** | **13.5** | N |
 
-img = Image.open('image.jpg').convert('RGB').resize((518, 518), Image.LANCZOS)
-data = np.expand_dims(np.array(img).astype(np.float32).transpose(2,0,1)/255.0, 0)
+> CoreML FP16 excluded due to poor accuracy (cos=0.944, SNR=9.5dB). CoreML's FP16 path uses half-precision for all intermediate activations, causing error accumulation in this deep ViT.
 
-descriptor = sess.run(None, {'images': data})[0]  # [B, 8448]
-```
+### ONNX Multi-Batch Throughput
 
-### CoreML (macOS only)
+| Model | B=1 | B=2 | B=4 | B=8 |
+|-------|-----|-----|-----|-----|
+| ONNX FP32 | 2.7 img/s | 2.3 img/s | 2.7 img/s | 2.9 img/s |
+| ONNX FP16 | 2.9 img/s | 2.6 img/s | 2.7 img/s | 2.9 img/s |
+| ONNX INT8 | 2.5 img/s | 2.1 img/s | 2.3 img/s | 2.5 img/s |
 
-```python
-import coremltools as ct, numpy as np
-from PIL import Image
+### Recommendations
+
+- **Best accuracy/size**: CoreML INT8 — 219 MB, 74ms, cos=0.9987, 13.5 FPS (macOS only)
+- **Cross-platform**: ONNX FP16 — 437 MB, 345ms, cos=0.99999, multi-batch
+- **Maximum compression**: ONNX INT8 — 223 MB, all platforms, multi-batch
+
+### On sub-8-bit quantization
+
+4-bit/1-bit quantization fails catastrophically (cos_sim ≈ 0). DINOv2 ViT has 12 blocks + aggregator (~200 weight tensors). Error from only 16 levels accumulates across 48+ MatMul layers. ViT architectures below 8-bit require **quantization-aware training**.
 
-model = ct.models.MLModel('megaloc_coreml_fp16.mlpackage')
+### Usage
 
-img = Image.open('image.jpg').convert('RGB').resize((518, 518), Image.LANCZOS)
+```bash
+uv sync
+
+# ONNX (cross-platform, multi-batch)
+python3 -c "
+import onnxruntime as ort, numpy as np
+from PIL import Image
+sess = ort.InferenceSession('megaloc_fp16.onnx', providers=['CPUExecutionProvider'])
+img = Image.open('img.jpg').convert('RGB').resize((518,518), Image.LANCZOS)
+data = np.expand_dims(np.array(img).astype(np.float32).transpose(2,0,1)/255.0, 0)
+desc = sess.run(None, {'images': data})[0]  # [B, 8448]
+"
+
+# CoreML (macOS, 5x faster)
+python3 -c "
+import coremltools as ct, numpy as np
+from PIL import Image
+model = ct.models.MLModel('megaloc_coreml_int8.mlpackage')
+img = Image.open('img.jpg').convert('RGB').resize((518,518), Image.LANCZOS)
 data = np.expand_dims(np.array(img).astype(np.float32).transpose(2,0,1)/255.0, 0)
+desc = list(model.predict({'images': data}).values())[0]
+"
+```
 
-descriptor = list(model.predict({'images': data}).values())[0]  # [B, 8448]
+### Reproduce
+
+```bash
+uv run python export_onnx_batch.py                             # ONNX multi-batch
+uv run python quantize_megaloc.py --method 8bit --input megaloc_fp32.onnx  # ONNX INT8
+uv run python convert_coreml.py                                # CoreML FP32/FP16/INT8
 ```
 
-## Reference
+---
+
+## Acknowledgements / Cite / BibTex
 
+If you use this repository please cite the following
 ```bibtex
 @inproceedings{Berton_2025_MegaLoc,
     author    = {Berton, Gabriele and Masone, Carlo},
     title     = {MegaLoc: One Retrieval to Place Them All},
-    booktitle = {CVPR Workshops},
+    booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops},
+    month     = {June},
     year      = {2025},
+    pages     = {2886--2892}
 }
 ```
+