commit repo mirror

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,171 @@
+---
+license: apache-2.0
+library_name: superpoint-pruned
+tags:
+  - SuperPoint
+---
+
+# 1. SuperPoint Pruned
+This repository focuses on optimizing the SuperPoint (SP) keypoint detection model for the TensorRT runtime on Jetson Orin Nano devices.
+To achieve better performance, we use methods like structural pruning, distillation training for performance recovery, and general model improvements for serialized model exporting.
+Despite the focus on the Jetson devices, the optimization methods are general enough to likely benefit other runtime backends as well and the ideas used can be applied to other similar network architectures.
+
+As a base we use the open source PyTorch implementation of Superpoint and corresponding base weights by Rémi Pautrat and Paul-Edouard Sarlin, which can be found [in this repository](https://github.com/rpautrat/SuperPoint).
+For distillation and evaluation, we use the [UZH FPV Dataset](https://fpv.ifi.uzh.ch/) [1].
+
+Note that there are also other popular SuperPoint implementations available, especially versions paired with keypoint matching models such as LightGlue. 
+In this repository we won't be using them because of their restrictive license and instead our work is intended more as a proof of concept to demonstrate feasibility of optimizing SuperPoint-like models. One can in theory apply the methods covered here to alternative implementations by making necessary structural changes and running the distillation pipeline with desired configurations.
+
+
+[1] J. Delmerico, T. Cieslewski, H. Rebecq, M. Faessler, and D. Scaramuzza,
+“Are We Ready for Autonomous Drone Racing? The UZH-FPV Drone Racing Dataset,”
+in *IEEE International Conference on Robotics and Automation (ICRA)*, 2019.
+
+# 2. Optimization methods
+
+The following provides a brief overview of the optimization methods used in the repository.
+
+## Pruning
+Pruning is the main optimization method used and main motivation for this repository. 
+For context, the network backbone of the default SuperPoint model architecture consists of eight sequential convolutional layers as shown in the figure:
+
+![superpoint_structure](./static/model_structure.png)
+
+The main observation is that after every two convolutional layers, the model has a pooling layer, which halves the width and height of the feature map.
+This means that the biggest bottlenecks of the model are the first convolutional layers, since they have to process the image at full input resolution.
+Hence the runtime can be improved by reducing the input/output channel sizes of desired convolutional layers, which our pruning interface allows to do with a simple class method addition.
+Note that in case of channel pruning, as a default we keep the channels that correspond to kernels with the highest mean absolute weight values, which is a good heuristic to retain as much of original performance as possible and also offers a good initialization point for potential distillation training.
+
+In the benchmarking section we outline two pruning configurations.
+One prunes only the `backbone_0_1` layer input channels from 64 to 32 (which means also adjusting `backbone_0_0` output channels accordingly) with no additional training, since this is the most expensive layer in the model.
+Another config we explore adjusts the first six layers of the model and because of that it requires distillation training to recover accuracy. 
+We provide the training checkpoint under `src/superpoint_pruning/weights/16_16_24_32_64.ckpt` and we use a combination of KL divergence and cross entropy loss for recovering the keypoint locations and a cosine similarity based loss for the descriptors.
+Here the filename `16_16_24_32_64` refers to the new values of input channels for the corresponding layers starting from `backbone_0_1`, `backbone_1_0` and so on.
+
+## Hierarchical top-k keypoint selection
+The original model extracts the final *k* keypoints by choosing the top scoring ones out of all possible pixel locations of the input image.
+Since the *k* value is normally much smaller than the number of candidate, then this can be done hierarchically by first splitting the image into chunks and finding the top *k* keypoints in each chunk, after which the final results can be aggregated.
+Using the same *k* and assuming no scoring ties, we are guaranteed the original result while getting a small speedup which is noticeable for constrained devices.
+
+## Skipping keypoint refinement
+
+This is not a direct optimization, but rather a useful observation in case the model runtime is especially important.
+The base model uses non-maximal supression to find top keypoints in each image region.
+While doing this, the original implementation also does two additional refinement steps, which allows to extract other high scoring keypoints that were supressed by the first pass.
+This usually results in more fine-grained keypoints in the final output, but in total the model does four extra passes over the full image resolution, which is very expensive.
+Hence we include an option to skip this refinement, as the outputs still stay relatively sensible without it, but this depends on the use case.
+
+# 3. Benchmark results
+In the following we compare the quality and inference time of five different model variations:
+- `original` - base SP with original weights
+- `original-topk` - base SP with hierarchical top-k enabled. Each variant uses the original *k* for sub-chunks, which guarantees the original results. The resolution 640x480 uses 32 chunks and 1920x1080 uses 36 chunks.
+- `pruned-light` - base SP where input channels of the `backbone_0_1` layer have been pruned from 64 to 32 with no additional training. This demonstrates the effect of the most expensive layer and how the kernel selection heuristic can retain performance. Previous top-k optimization is also enabled.
+- `pruned` - uses the `16_16_24_32_64` pruning config that affects the first six layers of the model (as described in the methods section) and uses our distillation checkpoint weights, which is trained on 250 examples and recovers the performance. Exact training configuration is referenced in the usage section. Previous top-k optimization is also enabled.
+- `pruned-noref` - previous `pruned` model with no NMS refinement steps to demonstrate the effects of it.
+
+## Quality evaluation
+
+The previous models were evaluated using PyTorch using the native dataset image size of 640x480 as input and 1024 keypoints in each case. We differentiate two specific dataset traces:
+- Indoor - 767 images from [UZH FPV Indoor forward facing trace #3](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/indoor_forward_3_snapdragon_with_gt.zip). This was the trace for training the `pruned` model and while we removed the direct images that were used for training, these results are still biased because of the sequential nature of the video frames.
+- Outdoor - 500 images from the [UZH FPV Outdoor forward facing trace #1](http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/outdoor_forward_1_snapdragon_with_gt.zip). Images from this trace were not used for performance recovery training of the pruned models and should serve as an example for the generalization capability of the distillation.
+
+We report two scores for each combination in the table below as *average number of keypoints covered / mean descriptor difference L2 norm*.
+The average number of keypoints covered score splits the original image into 8x8 chunks and checks if the target model predicted the same amount of keypoints in each chunk as the original. This isn't a perfect score but it also isn't as restrictive as exact match comparison and gives a general proxy score to assess, if the target model's keypoints have similar image coverage to the original (with 1024 being the best results, covering all original regions). In this case, a lower score doesn't mean that keypoints are definitely bad, but they might be more concentrated to a specific region than the original.
+The L2 norm is calculated by comparing the descriptor vectors at shared keypoint locations and calculating the norm of their difference. This score serves just as a quantitative signal for relative performance - in an ideal case this could be combined with a matching model to measure if the target model descriptor vectors are still recognizable.
+
+Note that generally, the best way to assess the pruned versions of SP is to test it in the intended downstream task and see if it still retains performance. We don't explore it here since we are not looking at a specific task currently.
+
+| Trace | original | original-topk | pruned-light | pruned | pruned-noref |
+|---|---:|---:|---:|---:|---:|
+| Inside | 1024 / 0 | 1024 / 0 | 637 / 0,92 | 817 / 0,24 | 700 / - |
+| Outside | 1024 / 0 | 1024 / 0 | 485 / 0,89 | 754 / 0,27 | 671 / - |
+
+In the following you can also see a few visualized examples of the predicted keypoints. The used models from left to right are `pruned-light`, `pruned` and `pruned-noref` respectively. The first set uses `image_1_1200.png` from the Indoor trace and the second set uses `image_0_1300.png` from the outdoor trace.
+
+All of the target model keypoints are compared to the original model predictions, hence the green points are keypoints which both models predicted exactly, while red and blue are points which one predicted, but the other did not.
+Notice that while there aren't many exactly overlapping points, there are many points where the original and target model predictions lie very close. Additionally, while not all of the target model prediction match the original, then often the predicted points are still sensible keypoint locations.
+
+| `pruned-light`|`pruned` |`pruned-noref` |
+|---|---|---|
+| ![](./static/pruned_light_indoor.png) | ![](./static/pruned_indoor.png) | ![](./static/pruned_noref_indoor.png) |
+
+| `pruned-light`|`pruned` |`pruned-noref` |
+|---|---|---|
+| ![](./static/pruned_light_outdoor.png) | ![](./static/pruned_outdoor.png) | ![](./static/pruned_noref_outdoor.png) |
+
+Alternatively, we can also compare the results with no overlay, e.g., here is the indoor image for the `pruned` model:
+
+![superpoint_structure](./static/pruned_indoor_separate.png)
+
+## Inference time evaluation
+For this section, all of the previous model configurations were exported as ONNX models and then serialized as TensorRT models using the `trtexec` CLI utility (with TRT version 10.3.0) with the `--fp16` flag to enable fp16 precision, which offers additional speedup for no quality degradation.
+Testing is done on a Jetson Orin Nano 8GB version and two times are reported: the first is the mean inference time across 1000 images from the previous indoor trace using our python benchmarking script and the second is the mean inference time reported by the TensorRT benchmarking utility (used as `trtexec --loadEngine=serialized_engine_file --useCudaGraph --useSpinWait`).
+All times are reported as milliseconds.
+
+| Input resolution | original | original-topk | pruned-light | pruned | pruned-noref |
+|---|---:|---:|---:|---:|---:|
+| 640 × 480 | 17,8 / 16,8 | 17,0 / 16,0 | 14,4 / 13,5 | 9,6 / 8,5 | 8,3 / 7,0 |
+| 1920 × 1080 | 111,1 / 109,6 | 102,5 / 102,1 | 86,9 / 84,6 | 54,2 / 52,9 | 45,1 / 43,5 |
+
+# 4. Project structure
+The code is organized into separate modules under `src/superpoint_pruning/`:
+- `export.py` - utility script for exporting a pruned SuperPoint model as ONNX.
+- `models/superpoint.py` - base optimized SuperPoint PyTorch model.
+- `distillation/` - PyTorch Lightning trainer script and utilities for distillation training. Includes `setup.py` script for downloading the base dataset that was used for training.
+- `evaluation/` - quality evaluation scripts and TensorRT wrapper for inference timing measurements.
+- `weights/` - base SuperPoint weights and precomputer distillation checkpoints.
+
+
+# 5. Setup and usage
+For installing the package with dependencies for running the basic distillation, evaluation, and export scripts, run:
+```
+pip install -r requirements.txt
+pip install -e .
+```
+To install the base dataset images (UZH FPV Indoor forward facing trace #3) used for training to the default path `src/superpoint_pruning/distillation/data`, run:
+```
+superpoint-pruning setup
+```
+Alternatively, specify the desired trace download link with the `--download-url` flag, e.g., the one for the UZH FPV Outdoor forward facing trace #1, and the local dataset directory name with `--output-dir-name`.
+
+## Run distillation training
+The base PyTorch Lightning training config that we used to produce the `pruned` model can be found under `src/superpoint_pruning/distillation/base_config.yaml`.
+Since this includes the hyperparameters and the data setup used for our training, then this should allow to reproduce the trained checkpoint that is provided under `src/superpoint_pruning/weights/16_16_24_32_64.ckpt`. The training script itself is very basic but can serve as a starting point for further experiments.
+To replicate the training, assuming that the main Indoor data trace is installed in the previous, run:
+```
+superpoint-pruning setup --generate-gt
+superpoint-pruning train
+```
+The first command computes and presaves the target keypoint and descriptor feature maps and the other initiate the Lightning training script.
+
+## Run evaluation
+The first row of quality evaluation results can be reproduced using the following commands corresponding to the different target models:
+```
+superpoint-pruning evaluate
+superpoint-pruning evaluate --hierarchical
+superpoint-pruning evaluate --backbone_0_1 32 --hierarchical
+superpoint-pruning evaluate --pruning-config ./src/superpoint_pruning/weights/16_16_24_32_64.ckpt --hierarchical
+superpoint-pruning evaluate --pruning-config ./src/superpoint_pruning/weights/16_16_24_32_64.ckpt --hierarchical --skip-refinement
+```
+Note that in this script and other similar ones, we can define the pruning either with format `--backbone_x_y new_input_channel_size` or by pointing it to a ckpt file named in the format `a_b_...ckpt`, which is translated to `--backbone_0_1 a`, `--backbone_1_0 b` and so on.
+For the evaluation using the outdoor trace (or some other trace), the `--image-dir` parameter has to be specified including the other specific options to use the same image range (`--no-skip --start-idx 1000 --end-idx 1500`).
+
+## Plot keypoint comparisons
+For a quick keypoint prediction comparsions plot between a target model and the original on a specific image, one can for example use:
+```
+superpoint-pruning plot --image-name image_1_1200.png --output ./test.png --num-keypoints 512 --pruning-config ./src/superpoint_pruning/weights/16_16_24_32_64.ckpt --overlay
+```
+This once again defaults to the UZH FPV Indoor forward facing trace #3 and accepts file names from the dataset.
+
+## Export model to ONNX
+To export the desired PyTorch model to ONNX:
+```
+superpoint-pruning export --backbone_0_1 32
+```
+The model pruning configuration can once again be controlled with same flags as shown in the evaluation example.
+
+## Run TRT inference benchmark
+In case the device supports TensorRT (with the necessary packages installed) and there is a serialized TRT model (here a dummy example of `superpoint.engine`), then the local python benchmark with a TRT wrapper can be run as:
+```
+superpoint-pruning benchmark --model-path superpoint.engine
+```

pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "superpoint-pruning"
+version = "0.1.0"
+description = "Pruning and distillation for SuperPoint"
+requires-python = ">=3.10"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[project.scripts]
+superpoint-pruning = "superpoint_pruning.cli:main"

requirements.txt

@@ -0,0 +1,11 @@
+# Base dependencies for export, distillation, evaluation, and plotting.
+# Install with: pip install -r requirements.txt
+# On Jetson, install NVIDIA's PyTorch first and skip the torch line if it is already present.
+numpy
+tqdm
+PyYAML
+matplotlib
+opencv-python
+lightning
+onnx
+torch==2.8.0

src/superpoint_pruning/cli.py

@@ -0,0 +1,46 @@
+import argparse
+
+from superpoint_pruning.evaluation import eval as eval_mod
+from superpoint_pruning import export as export_mod
+from superpoint_pruning.distillation import setup as setup_mod
+from superpoint_pruning.distillation import lightning_trainer as train_mod
+from superpoint_pruning.evaluation import benchmark_sp_tensorrt as benchmark_mod
+from superpoint_pruning.evaluation import plot_keypoints as plot_keypoints_mod
+
+
+def main():
+    parser = argparse.ArgumentParser(prog="superpoint-pruning")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    eval_parser = subparsers.add_parser("evaluate", help="Run SuperPoint evaluation")
+    eval_mod.add_parser_args(eval_parser)
+    eval_parser.set_defaults(func=eval_mod.main)
+
+    export_parser = subparsers.add_parser("export", help="Export SuperPoint to ONNX")
+    export_mod.add_parser_args(export_parser)
+    export_parser.set_defaults(func=export_mod.main)
+
+    setup_parser = subparsers.add_parser(
+        "setup", help="Setup the dataset for distillation and evaluation"
+    )
+    setup_mod.add_parser_args(setup_parser)
+    setup_parser.set_defaults(func=setup_mod.main)
+
+    train_parser = subparsers.add_parser("train", help="Train a pruned SuperPoint")
+    train_mod.add_parser_args(train_parser)
+    train_parser.set_defaults(func=train_mod.main)
+
+    benchmark_parser = subparsers.add_parser(
+        "benchmark", help="Benchmark a TensorRT SuperPoint engine"
+    )
+    benchmark_mod.add_parser_args(benchmark_parser)
+    benchmark_parser.set_defaults(func=benchmark_mod.main)
+
+    plot_parser = subparsers.add_parser(
+        "plot", help="Plot original and pruned SuperPoint keypoints"
+    )
+    plot_keypoints_mod.add_parser_args(plot_parser)
+    plot_parser.set_defaults(func=plot_keypoints_mod.main)
+
+    args = parser.parse_args()
+    args.func(args)

src/superpoint_pruning/distillation/base_config.yaml

@@ -0,0 +1,42 @@
+seed: 42
+
+model:
+  num_keypoints: 512
+
+  prune:
+    backbone_0_1: 16
+    backbone_1_0: 16
+    backbone_1_1: 24
+    backbone_2_0: 32
+    backbone_2_1: 64
+  
+data:
+  train_image_dir: "data/datasets/indoor_forward_3_snapdragon_with_gt/img"
+  train_image_ids_file: "data/gt_id_file_train_base_250.txt"
+  val_image_dir: "data/datasets/indoor_forward_3_snapdragon_with_gt/img"
+  val_image_ids_file: "data/gt_id_file_train_base_250.txt"
+  ground_truth_dir: "data"
+  keypoints_file: "kpts_train_base_250.npy"
+  descriptors_file: "desc_train_base_250.npy"
+  array_ids_file: "gt_id_file_train_base_250.txt"
+  image_size: [640, 480]
+  batch_size: 16
+  num_workers: 2
+  pin_memory: true
+
+optimizer:
+  lr: 1.0e-3
+  weight_decay: 1.0e-3
+
+loss:
+  cross_entropy_coef: 2.0
+  kl_coef: 2.0
+  descriptor_coef: 1.0
+
+trainer:
+  max_epochs: 300
+  accelerator: "auto"
+  devices: 1
+  precision: "32-true"
+  log_every_n_steps: 5
+  default_root_dir: "./distillation_runs"

src/superpoint_pruning/distillation/data/train_base.txt

@@ -0,0 +1,1575 @@
+image_1_1752.png
+image_0_1682.png
+image_0_1303.png
+image_1_1541.png
+image_0_1102.png
+image_0_1554.png
+image_1_1299.png
+image_1_1721.png
+image_0_1856.png
+image_1_1924.png
+image_0_1729.png
+image_0_1392.png
+image_1_1116.png
+image_0_1387.png
+image_0_1923.png
+image_0_1859.png
+image_1_1071.png
+image_0_1041.png
+image_1_1128.png
+image_1_1181.png
+image_1_1845.png
+image_1_1847.png
+image_0_1038.png
+image_0_1939.png
+image_1_1427.png
+image_1_1588.png
+image_1_1169.png
+image_1_1568.png
+image_0_1919.png
+image_1_1013.png
+image_0_1587.png
+image_0_1386.png
+image_0_1943.png
+image_1_1274.png
+image_0_1920.png
+image_0_1139.png
+image_1_1892.png
+image_0_1997.png
+image_1_1261.png
+image_0_1022.png
+image_0_1214.png
+image_0_1132.png
+image_0_1495.png
+image_1_1108.png
+image_1_1756.png
+image_0_1750.png
+image_0_1963.png
+image_1_1761.png
+image_1_1315.png
+image_0_1959.png
+image_1_1722.png
+image_0_1929.png
+image_0_1895.png
+image_0_1164.png
+image_0_1331.png
+image_0_1532.png
+image_0_1144.png
+image_1_1186.png
+image_0_1466.png
+image_0_1055.png
+image_1_1867.png
+image_0_1546.png
+image_0_1595.png
+image_1_1888.png
+image_1_1611.png
+image_1_1615.png
+image_1_1848.png
+image_0_1179.png
+image_0_1675.png
+image_0_1460.png
+image_0_1142.png
+image_1_1868.png
+image_1_1309.png
+image_1_1237.png
+image_0_1236.png
+image_0_1215.png
+image_0_1123.png
+image_1_1529.png
+image_0_1590.png
+image_0_1057.png
+image_1_1006.png
+image_1_1928.png
+image_1_1776.png
+image_1_1757.png
+image_0_1044.png
+image_0_1036.png
+image_1_1294.png
+image_1_1100.png
+image_0_1950.png
+image_1_1194.png
+image_1_1968.png
+image_0_1027.png
+image_0_1643.png
+image_0_1527.png
+image_0_1207.png
+image_1_1346.png
+image_0_1660.png
+image_0_1228.png
+image_1_1417.png
+image_0_1368.png
+image_1_1523.png
+image_0_1421.png
+image_0_1972.png
+image_1_1462.png
+image_0_1016.png
+image_0_1630.png
+image_1_1911.png
+image_1_1954.png
+image_0_1948.png
+image_0_1809.png
+image_0_1822.png
+image_1_1491.png
+image_1_1022.png
+image_0_1621.png
+image_0_1362.png
+image_1_1745.png
+image_0_1081.png
+image_0_1562.png
+image_1_1790.png
+image_0_1909.png
+image_1_1494.png
+image_0_1378.png
+image_0_1450.png
+image_1_1146.png
+image_0_1459.png
+image_0_1503.png
+image_1_1688.png
+image_0_1600.png
+image_0_1089.png
+image_1_1880.png
+image_0_1879.png
+image_0_1379.png
+image_1_1540.png
+image_0_1115.png
+image_0_1741.png
+image_0_1861.png
+image_0_1581.png
+image_1_1645.png
+image_0_1867.png
+image_1_1048.png
+image_0_1275.png
+image_0_1073.png
+image_0_1912.png
+image_0_1565.png
+image_0_1500.png
+image_0_1603.png
+image_1_1479.png
+image_0_1656.png
+image_0_1582.png
+image_1_1425.png
+image_1_1003.png
+image_1_1130.png
+image_0_1003.png
+image_1_1148.png
+image_1_1905.png
+image_1_1641.png
+image_1_1599.png
+image_0_1956.png
+image_0_1286.png
+image_1_1104.png
+image_1_1827.png
+image_1_1455.png
+image_1_1778.png
+image_0_1588.png
+image_1_1875.png
+image_0_1395.png
+image_0_1760.png
+image_0_1677.png
+image_0_1188.png
+image_0_1700.png
+image_0_1672.png
+image_1_1957.png
+image_0_1993.png
+image_1_1987.png
+image_1_1706.png
+image_1_1421.png
+image_1_1978.png
+image_0_1357.png
+image_1_1591.png
+image_0_1086.png
+image_0_1464.png
+image_0_1699.png
+image_1_1434.png
+image_1_1857.png
+image_0_1978.png
+image_1_1247.png
+image_0_1611.png
+image_0_1278.png
+image_0_1502.png
+image_1_1578.png
+image_1_1530.png
+image_0_1735.png
+image_1_1962.png
+image_1_1903.png
+image_0_1725.png
+image_0_1279.png
+image_1_1582.png
+image_1_1040.png
+image_1_1946.png
+image_0_1320.png
+image_0_1238.png
+image_1_1550.png
+image_1_1808.png
+image_0_1727.png
+image_0_1190.png
+image_0_1348.png
+image_1_1930.png
+image_1_1948.png
+image_1_1449.png
+image_1_1974.png
+image_1_1670.png
+image_1_1552.png
+image_0_1522.png
+image_0_1973.png
+image_0_1734.png
+image_0_1796.png
+image_0_1801.png
+image_1_1564.png
+image_1_1236.png
+image_0_1471.png
+image_0_1468.png
+image_0_1960.png
+image_0_1807.png
+image_0_1833.png
+image_0_1905.png
+image_0_1768.png
+image_1_1750.png
+image_0_1137.png
+image_1_1230.png
+image_0_1114.png
+image_1_1769.png
+image_1_1686.png
+image_0_1441.png
+image_1_1956.png
+image_1_1458.png
+image_1_1316.png
+image_1_1809.png
+image_0_1052.png
+image_0_1515.png
+image_0_1152.png
+image_0_1918.png
+image_0_1337.png
+image_1_1807.png
+image_0_1767.png
+image_1_1595.png
+image_1_1637.png
+image_1_1823.png
+image_0_1283.png
+image_1_1073.png
+image_0_1366.png
+image_0_1370.png
+image_0_1186.png
+image_1_1514.png
+image_1_1337.png
+image_0_1234.png
+image_1_1942.png
+image_1_1597.png
+image_1_1303.png
+image_0_1393.png
+image_0_1947.png
+image_1_1679.png
+image_1_1354.png
+image_0_1196.png
+image_0_1195.png
+image_1_1245.png
+image_0_1406.png
+image_0_1116.png
+image_1_1664.png
+image_0_1162.png
+image_1_1152.png
+image_0_1622.png
+image_1_1310.png
+image_0_1457.png
+image_0_1074.png
+image_0_1451.png
+image_1_1105.png
+image_1_1051.png
+image_1_1803.png
+image_1_1037.png
+image_1_1830.png
+image_0_1145.png
+image_0_1541.png
+image_1_1557.png
+image_1_1285.png
+image_1_1587.png
+image_1_1437.png
+image_1_1559.png
+image_1_1306.png
+image_0_1167.png
+image_1_1380.png
+image_0_1284.png
+image_0_1006.png
+image_1_1496.png
+image_0_1641.png
+image_0_1658.png
+image_1_1047.png
+image_1_1370.png
+image_1_1155.png
+image_1_1453.png
+image_1_1257.png
+image_1_1359.png
+image_0_1812.png
+image_0_1063.png
+image_1_1573.png
+image_0_1092.png
+image_0_1248.png
+image_0_1509.png
+image_0_1001.png
+image_1_1286.png
+image_1_1166.png
+image_0_1550.png
+image_0_1628.png
+image_1_1382.png
+image_1_1439.png
+image_0_1151.png
+image_1_1785.png
+image_0_1635.png
+image_0_1780.png
+image_1_1433.png
+image_1_1086.png
+image_1_1558.png
+image_1_1210.png
+image_0_1753.png
+image_0_1563.png
+image_0_1689.png
+image_0_1813.png
+image_1_1387.png
+image_1_1915.png
+image_1_1579.png
+image_1_1206.png
+image_0_1002.png
+image_0_1992.png
+image_0_1229.png
+image_1_1394.png
+image_0_1795.png
+image_1_1565.png
+image_1_1406.png
+image_1_1445.png
+image_0_1013.png
+image_0_1280.png
+image_0_1783.png
+image_1_1897.png
+image_1_1604.png
+image_1_1996.png
+image_1_1632.png
+image_0_1521.png
+image_1_1327.png
+image_0_1764.png
+image_1_1068.png
+image_1_1214.png
+image_0_1798.png
+image_0_1819.png
+image_0_1434.png
+image_0_1511.png
+image_1_1959.png
+image_0_1070.png
+image_0_1692.png
+image_0_1815.png
+image_0_1315.png
+image_1_1521.png
+image_1_1717.png
+image_0_1033.png
+image_0_1951.png
+image_1_1952.png
+image_0_1203.png
+image_0_1175.png
+image_1_1876.png
+image_1_1649.png
+image_0_1523.png
+image_0_1720.png
+image_1_1793.png
+image_0_1419.png
+image_0_1994.png
+image_1_1739.png
+image_0_1520.png
+image_1_1918.png
+image_0_1719.png
+image_1_1730.png
+image_1_1275.png
+image_1_1012.png
+image_0_1223.png
+image_1_1674.png
+image_1_1883.png
+image_0_1166.png
+image_1_1603.png
+image_0_1383.png
+image_1_1312.png
+image_0_1129.png
+image_1_1289.png
+image_1_1250.png
+image_1_1628.png
+image_0_1226.png
+image_1_1889.png
+image_1_1443.png
+image_1_1376.png
+image_1_1096.png
+image_0_1721.png
+image_0_1031.png
+image_0_1804.png
+image_0_1717.png
+image_0_1130.png
+image_0_1642.png
+image_1_1834.png
+image_0_1980.png
+image_0_1707.png
+image_1_1358.png
+image_0_1192.png
+image_1_1546.png
+image_1_1334.png
+image_0_1695.png
+image_1_1467.png
+image_1_1478.png
+image_1_1690.png
+image_0_1090.png
+image_0_1954.png
+image_0_1099.png
+image_1_1383.png
+image_1_1280.png
+image_0_1825.png
+image_1_1060.png
+image_1_1544.png
+image_1_1507.png
+image_0_1967.png
+image_0_1619.png
+image_1_1967.png
+image_0_1479.png
+image_0_1270.png
+image_0_1847.png
+image_0_1638.png
+image_1_1940.png
+image_1_1574.png
+image_1_1298.png
+image_1_1045.png
+image_0_1350.png
+image_1_1080.png
+image_1_1511.png
+image_1_1137.png
+image_1_1601.png
+image_0_1542.png
+image_1_1636.png
+image_0_1217.png
+image_0_1206.png
+image_1_1758.png
+image_1_1995.png
+image_1_1147.png
+image_1_1972.png
+image_1_1392.png
+image_0_1917.png
+image_1_1869.png
+image_0_1744.png
+image_0_1811.png
+image_0_1359.png
+image_1_1618.png
+image_0_1043.png
+image_1_1542.png
+image_0_1210.png
+image_1_1476.png
+image_0_1854.png
+image_0_1456.png
+image_0_1014.png
+image_0_1623.png
+image_0_1892.png
+image_0_1932.png
+image_0_1497.png
+image_1_1667.png
+image_1_1567.png
+image_0_1470.png
+image_1_1084.png
+image_0_1824.png
+image_1_1420.png
+image_0_1962.png
+image_0_1463.png
+image_0_1539.png
+image_1_1788.png
+image_1_1894.png
+image_1_1454.png
+image_0_1032.png
+image_0_1551.png
+image_1_1348.png
+image_0_1173.png
+image_0_1181.png
+image_0_1836.png
+image_1_1413.png
+image_1_1362.png
+image_0_1608.png
+image_0_1657.png
+image_0_1606.png
+image_1_1719.png
+image_0_1525.png
+image_0_1382.png
+image_1_1654.png
+image_1_1512.png
+image_1_1833.png
+image_0_1417.png
+image_1_1818.png
+image_0_1204.png
+image_1_1920.png
+image_1_1533.png
+image_1_1561.png
+image_1_1222.png
+image_1_1816.png
+image_1_1555.png
+image_0_1483.png
+image_1_1492.png
+image_0_1472.png
+image_0_1916.png
+image_0_1762.png
+image_1_1017.png
+image_0_1792.png
+image_1_1242.png
+image_0_1852.png
+image_1_1372.png
+image_0_1530.png
+image_1_1630.png
+image_0_1375.png
+image_1_1772.png
+image_0_1373.png
+image_1_1945.png
+image_1_1749.png
+image_0_1096.png
+image_0_1263.png
+image_0_1965.png
+image_1_1246.png
+image_1_1910.png
+image_0_1281.png
+image_0_1076.png
+image_0_1494.png
+image_0_1816.png
+image_0_1868.png
+image_1_1356.png
+image_1_1057.png
+image_1_1302.png
+image_1_1267.png
+image_0_1938.png
+image_0_1633.png
+image_1_1142.png
+image_0_1258.png
+image_0_1491.png
+image_0_1851.png
+image_0_1821.png
+image_0_1333.png
+image_0_1747.png
+image_0_1991.png
+image_1_1851.png
+image_1_1094.png
+image_0_1399.png
+image_0_1059.png
+image_1_1343.png
+image_0_1156.png
+image_0_1713.png
+image_0_1416.png
+image_0_1714.png
+image_0_1952.png
+image_1_1989.png
+image_0_1244.png
+image_1_1701.png
+image_0_1300.png
+image_0_1647.png
+image_1_1446.png
+image_0_1165.png
+image_1_1682.png
+image_1_1646.png
+image_1_1125.png
+image_0_1285.png
+image_1_1196.png
+image_1_1866.png
+image_0_1047.png
+image_1_1231.png
+image_0_1135.png
+image_1_1087.png
+image_1_1500.png
+image_1_1357.png
+image_0_1893.png
+image_0_1071.png
+image_0_1191.png
+image_1_1039.png
+image_1_1639.png
+image_0_1449.png
+image_0_1155.png
+image_0_1061.png
+image_1_1426.png
+image_0_1222.png
+image_1_1190.png
+image_0_1789.png
+image_0_1941.png
+image_0_1988.png
+image_1_1062.png
+image_1_1079.png
+image_1_1290.png
+image_0_1591.png
+image_0_1908.png
+image_1_1694.png
+image_1_1650.png
+image_1_1122.png
+image_1_1014.png
+image_0_1696.png
+image_1_1183.png
+image_0_1394.png
+image_1_1342.png
+image_0_1018.png
+image_0_1087.png
+image_0_1037.png
+image_1_1070.png
+image_1_1384.png
+image_1_1218.png
+image_0_1189.png
+image_0_1363.png
+image_0_1012.png
+image_1_1388.png
+image_1_1058.png
+image_1_1715.png
+image_1_1410.png
+image_0_1853.png
+image_0_1589.png
+image_0_1496.png
+image_1_1678.png
+image_0_1930.png
+image_1_1208.png
+image_1_1470.png
+image_0_1024.png
+image_0_1650.png
+image_0_1516.png
+image_0_1678.png
+image_0_1492.png
+image_1_1873.png
+image_0_1594.png
+image_1_1904.png
+image_1_1031.png
+image_1_1885.png
+image_0_1119.png
+image_0_1644.png
+image_0_1906.png
+image_1_1377.png
+image_1_1279.png
+image_1_1538.png
+image_1_1986.png
+image_1_1887.png
+image_1_1802.png
+image_1_1441.png
+image_0_1325.png
+image_1_1490.png
+image_0_1060.png
+image_0_1610.png
+image_1_1549.png
+image_1_1509.png
+image_0_1831.png
+image_0_1476.png
+image_1_1754.png
+image_1_1389.png
+image_0_1710.png
+image_1_1318.png
+image_1_1513.png
+image_1_1687.png
+image_0_1842.png
+image_1_1759.png
+image_0_1062.png
+image_1_1705.png
+image_1_1179.png
+image_0_1519.png
+image_0_1051.png
+image_0_1596.png
+image_1_1055.png
+image_1_1225.png
+image_0_1999.png
+image_0_1732.png
+image_1_1771.png
+image_1_1097.png
+image_1_1608.png
+image_0_1774.png
+image_0_1308.png
+image_0_1239.png
+image_0_1564.png
+image_0_1770.png
+image_1_1899.png
+image_0_1612.png
+image_0_1841.png
+image_0_1974.png
+image_1_1273.png
+image_1_1009.png
+image_1_1074.png
+image_1_1607.png
+image_0_1766.png
+image_0_1671.png
+image_1_1944.png
+image_1_1891.png
+image_0_1475.png
+image_0_1781.png
+image_1_1993.png
+image_0_1458.png
+image_1_1301.png
+image_1_1284.png
+image_1_1629.png
+image_1_1249.png
+image_1_1184.png
+image_1_1076.png
+image_0_1334.png
+image_0_1049.png
+image_1_1861.png
+image_0_1172.png
+image_0_1761.png
+image_0_1585.png
+image_1_1056.png
+image_0_1971.png
+image_0_1517.png
+image_1_1605.png
+image_0_1436.png
+image_0_1282.png
+image_0_1803.png
+image_1_1676.png
+image_1_1390.png
+image_1_1515.png
+image_1_1475.png
+image_1_1154.png
+image_0_1025.png
+image_1_1235.png
+image_1_1320.png
+image_1_1409.png
+image_1_1452.png
+image_1_1120.png
+image_0_1910.png
+image_0_1986.png
+image_1_1525.png
+image_1_1175.png
+image_0_1444.png
+image_1_1723.png
+image_0_1873.png
+image_0_1722.png
+image_0_1914.png
+image_1_1336.png
+image_0_1567.png
+image_0_1322.png
+image_1_1020.png
+image_1_1167.png
+image_1_1369.png
+image_0_1935.png
+image_1_1191.png
+image_0_1898.png
+image_1_1232.png
+image_0_1122.png
+image_1_1840.png
+image_0_1618.png
+image_0_1535.png
+image_1_1419.png
+image_1_1941.png
+image_0_1931.png
+image_1_1672.png
+image_0_1069.png
+image_0_1936.png
+image_1_1976.png
+image_1_1797.png
+image_1_1613.png
+image_1_1328.png
+image_1_1829.png
+image_1_1638.png
+image_1_1842.png
+image_0_1662.png
+image_1_1075.png
+image_0_1745.png
+image_1_1177.png
+image_0_1381.png
+image_1_1753.png
+image_0_1634.png
+image_0_1302.png
+image_0_1109.png
+image_1_1660.png
+image_1_1252.png
+image_1_1657.png
+image_1_1831.png
+image_0_1877.png
+image_0_1712.png
+image_1_1489.png
+image_1_1018.png
+image_1_1988.png
+image_0_1340.png
+image_1_1616.png
+image_0_1324.png
+image_1_1625.png
+image_0_1583.png
+image_1_1792.png
+image_0_1826.png
+image_1_1935.png
+image_0_1560.png
+image_0_1358.png
+image_0_1984.png
+image_1_1459.png
+image_0_1065.png
+image_0_1314.png
+image_1_1795.png
+image_1_1726.png
+image_1_1260.png
+image_1_1317.png
+image_1_1371.png
+image_1_1272.png
+image_0_1404.png
+image_1_1164.png
+image_0_1654.png
+image_1_1432.png
+image_0_1093.png
+image_0_1084.png
+image_1_1240.png
+image_0_1886.png
+image_1_1925.png
+image_1_1197.png
+image_0_1869.png
+image_1_1030.png
+image_0_1171.png
+image_0_1524.png
+image_1_1725.png
+image_0_1659.png
+image_1_1937.png
+image_1_1204.png
+image_1_1583.png
+image_1_1428.png
+image_1_1465.png
+image_1_1871.png
+image_1_1916.png
+image_1_1182.png
+image_0_1158.png
+image_0_1570.png
+image_0_1140.png
+image_1_1695.png
+image_1_1783.png
+image_1_1736.png
+image_0_1299.png
+image_1_1766.png
+image_0_1846.png
+image_1_1050.png
+image_1_1734.png
+image_0_1443.png
+image_1_1448.png
+image_0_1547.png
+image_0_1601.png
+image_1_1408.png
+image_0_1850.png
+image_1_1127.png
+image_1_1620.png
+image_0_1023.png
+image_0_1778.png
+image_0_1017.png
+image_0_1970.png
+image_0_1733.png
+image_1_1393.png
+image_0_1505.png
+image_1_1168.png
+image_1_1170.png
+image_1_1870.png
+image_0_1293.png
+image_0_1019.png
+image_1_1364.png
+image_0_1773.png
+image_0_1414.png
+image_1_1621.png
+image_1_1044.png
+image_1_1227.png
+image_0_1377.png
+image_0_1704.png
+image_1_1107.png
+image_1_1234.png
+image_0_1730.png
+image_1_1189.png
+image_0_1213.png
+image_0_1298.png
+image_0_1785.png
+image_1_1367.png
+image_1_1997.png
+image_0_1982.png
+image_0_1998.png
+image_1_1737.png
+image_1_1858.png
+image_0_1976.png
+image_1_1864.png
+image_1_1207.png
+image_0_1445.png
+image_1_1853.png
+image_1_1822.png
+image_1_1139.png
+image_1_1893.png
+image_0_1058.png
+image_1_1572.png
+image_1_1345.png
+image_0_1740.png
+image_0_1085.png
+image_0_1556.png
+image_1_1516.png
+image_1_1053.png
+image_0_1987.png
+image_0_1631.png
+image_0_1053.png
+image_0_1141.png
+image_1_1780.png
+image_0_1703.png
+image_1_1590.png
+image_1_1675.png
+image_1_1027.png
+image_0_1548.png
+image_0_1291.png
+image_1_1415.png
+image_0_1420.png
+image_0_1161.png
+image_0_1356.png
+image_0_1143.png
+image_1_1407.png
+image_0_1499.png
+image_0_1313.png
+image_1_1092.png
+image_0_1227.png
+image_0_1199.png
+image_1_1131.png
+image_0_1498.png
+image_1_1052.png
+image_1_1258.png
+image_1_1979.png
+image_0_1462.png
+image_1_1651.png
+image_1_1763.png
+image_1_1217.png
+image_0_1814.png
+image_0_1891.png
+image_0_1429.png
+image_0_1020.png
+image_1_1609.png
+image_1_1291.png
+image_1_1517.png
+image_0_1346.png
+image_1_1497.png
+image_0_1169.png
+image_0_1985.png
+image_1_1466.png
+image_0_1668.png
+image_1_1661.png
+image_0_1864.png
+image_0_1412.png
+image_1_1874.png
+image_0_1942.png
+image_0_1937.png
+image_0_1323.png
+image_1_1036.png
+image_1_1980.png
+image_1_1163.png
+image_1_1681.png
+image_1_1919.png
+image_1_1221.png
+image_0_1111.png
+image_1_1852.png
+image_0_1620.png
+image_1_1539.png
+image_0_1413.png
+image_0_1579.png
+image_0_1400.png
+image_0_1410.png
+image_0_1326.png
+image_1_1375.png
+image_0_1646.png
+image_1_1488.png
+image_0_1834.png
+image_1_1635.png
+image_0_1469.png
+image_1_1806.png
+image_0_1230.png
+image_0_1157.png
+image_1_1431.png
+image_0_1665.png
+image_0_1118.png
+image_0_1112.png
+image_0_1233.png
+image_1_1648.png
+image_1_1964.png
+image_1_1856.png
+image_0_1946.png
+image_1_1553.png
+image_0_1355.png
+image_1_1707.png
+image_0_1154.png
+image_1_1323.png
+image_0_1708.png
+image_0_1534.png
+image_1_1832.png
+image_0_1273.png
+image_1_1716.png
+image_0_1609.png
+image_0_1232.png
+image_1_1135.png
+image_0_1996.png
+image_1_1444.png
+image_0_1095.png
+image_1_1054.png
+image_1_1161.png
+image_1_1019.png
+image_1_1038.png
+image_0_1673.png
+image_1_1460.png
+image_1_1619.png
+image_1_1349.png
+image_0_1624.png
+image_0_1514.png
+image_1_1877.png
+image_0_1094.png
+image_0_1823.png
+image_0_1926.png
+image_0_1454.png
+image_1_1566.png
+image_0_1797.png
+image_0_1254.png
+image_1_1782.png
+image_1_1837.png
+image_1_1322.png
+image_1_1209.png
+image_0_1015.png
+image_0_1629.png
+image_1_1815.png
+image_0_1294.png
+image_0_1584.png
+image_0_1105.png
+image_0_1398.png
+image_1_1640.png
+image_0_1353.png
+image_1_1878.png
+image_1_1614.png
+image_0_1321.png
+image_0_1433.png
+image_1_1606.png
+image_0_1683.png
+image_1_1560.png
+image_0_1701.png
+image_0_1255.png
+image_0_1655.png
+image_1_1263.png
+image_1_1817.png
+image_0_1010.png
+image_0_1150.png
+image_1_1069.png
+image_0_1810.png
+image_0_1686.png
+image_0_1706.png
+image_0_1345.png
+image_1_1321.png
+image_0_1625.png
+image_0_1901.png
+image_1_1473.png
+image_1_1537.png
+image_1_1786.png
+image_1_1810.png
+image_1_1110.png
+image_1_1825.png
+image_0_1510.png
+image_0_1301.png
+image_1_1811.png
+image_0_1465.png
+image_0_1666.png
+image_1_1140.png
+image_1_1297.png
+image_0_1736.png
+image_1_1900.png
+image_0_1977.png
+image_1_1212.png
+image_1_1751.png
+image_1_1647.png
+image_0_1488.png
+image_1_1332.png
+image_0_1124.png
+image_1_1268.png
+image_1_1109.png
+image_1_1789.png
+image_0_1617.png
+image_1_1259.png
+image_1_1998.png
+image_1_1102.png
+image_1_1729.png
+image_0_1292.png
+image_1_1099.png
+image_0_1883.png
+image_0_1265.png
+image_0_1857.png
+image_0_1177.png
+image_0_1928.png
+image_1_1451.png
+image_0_1779.png
+image_0_1288.png
+image_1_1063.png
+image_0_1388.png
+image_1_1338.png
+image_0_1545.png
+image_1_1774.png
+image_0_1174.png
+image_0_1453.png
+image_0_1489.png
+image_1_1580.png
+image_1_1821.png
+image_1_1485.png
+image_1_1879.png
+image_0_1537.png
+image_1_1028.png
+image_0_1264.png
+image_0_1208.png
+image_1_1325.png
+image_0_1924.png
+image_1_1922.png
+image_0_1341.png
+image_1_1078.png
+image_1_1344.png
+image_1_1180.png
+image_0_1990.png
+image_1_1360.png
+image_0_1648.png
+image_0_1828.png
+image_0_1793.png
+image_0_1344.png
+image_0_1549.png
+image_1_1784.png
+image_0_1045.png
+image_0_1403.png
+image_1_1007.png
+image_0_1685.png
+image_0_1578.png
+image_1_1229.png
+image_0_1755.png
+image_1_1983.png
+image_1_1335.png
+image_1_1961.png
+image_0_1800.png
+image_0_1080.png
+image_1_1700.png
+image_1_1025.png
+image_1_1300.png
+image_0_1136.png
+image_0_1289.png
+image_1_1308.png
+image_0_1437.png
+image_0_1439.png
+image_1_1531.png
+image_0_1424.png
+image_1_1865.png
+image_0_1688.png
+image_0_1163.png
+image_1_1733.png
+image_1_1188.png
+image_0_1533.png
+image_1_1141.png
+image_1_1768.png
+image_0_1543.png
+image_1_1846.png
+image_0_1575.png
+image_0_1007.png
+image_0_1765.png
+image_0_1335.png
+image_0_1955.png
+image_1_1819.png
+image_1_1838.png
+image_0_1480.png
+image_1_1794.png
+image_0_1904.png
+image_0_1602.png
+image_0_1922.png
+image_1_1991.png
+image_1_1932.png
+image_0_1253.png
+image_1_1882.png
+image_0_1775.png
+image_1_1385.png
+image_1_1543.png
+image_0_1128.png
+image_0_1409.png
+image_1_1251.png
+image_0_1402.png
+image_0_1328.png
+image_0_1508.png
+image_0_1440.png
+image_0_1607.png
+image_1_1826.png
+image_0_1913.png
+image_0_1272.png
+image_1_1011.png
+image_0_1369.png
+image_1_1610.png
+image_1_1424.png
+image_0_1209.png
+image_0_1040.png
+image_0_1501.png
+image_1_1077.png
+image_1_1314.png
+image_1_1510.png
+image_1_1576.png
+image_0_1452.png
+image_1_1243.png
+image_0_1949.png
+image_1_1066.png
+image_0_1940.png
+image_1_1781.png
+image_0_1571.png
+image_0_1802.png
+image_1_1943.png
+image_0_1652.png
+image_1_1642.png
+image_0_1506.png
+image_0_1243.png
+image_0_1056.png
+image_1_1813.png
+image_1_1755.png
+image_0_1327.png
+image_1_1046.png
+image_1_1653.png
+image_0_1316.png
+image_1_1760.png
+image_1_1145.png
+image_0_1290.png
+image_1_1939.png
+image_1_1123.png
+image_1_1958.png
+image_1_1524.png
+image_1_1293.png
+image_0_1907.png
+image_1_1255.png
+image_1_1594.png
+image_1_1652.png
+image_1_1401.png
+image_0_1106.png
+image_0_1054.png
+image_1_1340.png
+image_1_1311.png
+image_1_1934.png
+image_0_1261.png
+image_1_1728.png
+image_1_1082.png
+image_1_1775.png
+image_0_1639.png
+image_0_1844.png
+image_1_1977.png
+image_0_1435.png
+image_0_1405.png
+image_1_1487.png
+image_0_1113.png
+image_1_1593.png
+image_0_1878.png
+image_1_1277.png
+image_1_1468.png
+image_0_1005.png
+image_0_1277.png
+image_1_1307.png
+image_0_1349.png
+image_1_1339.png
+image_0_1737.png
+image_1_1960.png
+image_1_1791.png
+image_0_1961.png
+image_1_1113.png
+image_0_1848.png
+image_1_1352.png
+image_1_1153.png
+image_0_1574.png
+image_0_1026.png
+image_1_1481.png
+image_0_1097.png
+image_0_1690.png
+image_1_1378.png
+image_1_1223.png
+image_1_1411.png
+image_1_1034.png
+image_0_1107.png
+image_0_1390.png
+image_1_1347.png
+image_0_1536.png
+image_1_1305.png
+image_1_1200.png
+image_1_1292.png
+image_1_1764.png
+image_1_1093.png
+image_1_1265.png
+image_1_1368.png
+image_0_1146.png
+image_0_1259.png
+image_0_1758.png
+image_1_1534.png
+image_0_1576.png
+image_1_1463.png
+image_1_1435.png
+image_0_1147.png
+image_1_1504.png
+image_0_1339.png
+image_0_1426.png
+image_1_1562.png
+image_0_1431.png
+image_1_1953.png
+image_0_1415.png
+image_1_1912.png
+image_1_1241.png
+image_1_1696.png
+image_1_1015.png
+image_1_1185.png
+image_1_1032.png
+image_0_1902.png
+image_0_1669.png
+image_1_1907.png
+image_0_1225.png
+image_1_1898.png
+image_1_1689.png
+image_1_1859.png
+image_0_1371.png
+image_1_1254.png
+image_1_1836.png
+image_1_1035.png
+image_0_1742.png
+image_0_1262.png
+image_0_1216.png
+image_0_1372.png
+image_1_1677.png
+image_1_1709.png
+image_0_1438.png
+image_1_1596.png
+image_1_1438.png
+image_0_1964.png
+image_1_1520.png
+image_0_1446.png
+image_1_1746.png
+image_0_1481.png
+image_1_1741.png
+image_1_1144.png
+image_0_1805.png
+image_0_1531.png
+image_0_1160.png
+image_0_1716.png
+image_1_1505.png
+image_1_1416.png
+image_0_1461.png
+image_0_1138.png
+image_1_1024.png
+image_0_1772.png
+image_0_1794.png
+image_1_1262.png
+image_1_1992.png
+image_0_1231.png
+image_1_1350.png
+image_1_1101.png
+image_1_1917.png
+image_0_1257.png
+image_0_1830.png
+image_1_1547.png
+image_0_1552.png
+image_1_1329.png
+image_1_1160.png
+image_0_1849.png
+image_0_1711.png
+image_1_1042.png
+image_0_1442.png
+image_1_1351.png
+image_1_1577.png
+image_1_1502.png
+image_0_1311.png
+image_0_1818.png
+image_0_1110.png
+image_0_1121.png
+image_0_1885.png
+image_1_1973.png
+image_0_1245.png
+image_1_1623.png
+image_0_1008.png
+image_0_1427.png
+image_0_1702.png
+image_1_1442.png
+image_0_1927.png
+image_0_1969.png
+image_1_1205.png
+image_0_1423.png
+image_0_1028.png
+image_1_1950.png
+image_0_1064.png
+image_0_1193.png
+image_0_1559.png
+image_1_1486.png
+image_0_1784.png
+image_1_1043.png
+image_1_1201.png
+image_1_1115.png
+image_1_1361.png
+image_0_1187.png
+image_1_1933.png
+image_1_1396.png
+image_0_1385.png
+image_0_1256.png
+image_0_1200.png
+image_0_1418.png
+image_0_1820.png
+image_1_1269.png
+image_1_1386.png
+image_1_1365.png
+image_0_1310.png
+image_1_1287.png
+image_0_1989.png
+image_0_1598.png
+image_1_1016.png
+image_0_1477.png
+image_1_1740.png
+image_0_1679.png
+image_1_1324.png
+image_0_1389.png
+image_1_1072.png
+image_1_1908.png
+image_0_1957.png
+image_0_1079.png
+image_0_1718.png
+image_1_1697.png
+image_0_1561.png
+image_0_1738.png
+image_1_1414.png
+image_1_1929.png
+image_0_1343.png
+image_1_1049.png
+image_1_1627.png
+image_0_1769.png
+image_1_1238.png
+image_0_1380.png
+image_0_1724.png
+image_1_1430.png
+image_0_1983.png
+image_0_1979.png
+image_0_1731.png
+image_1_1450.png
+image_1_1304.png
+image_0_1945.png
+image_1_1671.png
+image_1_1767.png
+image_1_1288.png
+image_1_1622.png
+image_1_1151.png
+image_0_1271.png
+image_0_1108.png
+image_0_1843.png
+image_0_1958.png
+image_1_1644.png
+image_1_1896.png
+image_1_1693.png
+image_1_1134.png
+image_1_1718.png
+image_1_1403.png
+image_0_1739.png
+image_1_1570.png
+image_0_1776.png
+image_0_1786.png
+image_1_1563.png
+image_1_1712.png
+image_1_1828.png
+image_1_1483.png
+image_1_1955.png
+image_1_1172.png
+image_0_1104.png
+image_0_1408.png
+image_0_1862.png
+image_1_1931.png
+image_0_1485.png
+image_1_1963.png
+image_0_1478.png
+image_0_1697.png
+image_1_1326.png
+image_1_1226.png
+image_1_1724.png
+image_1_1703.png
+image_0_1194.png
+image_0_1705.png
+image_1_1173.png
+image_0_1306.png
+image_0_1752.png
+image_0_1787.png
+image_1_1199.png
+image_0_1295.png
+image_0_1674.png
+image_0_1319.png
+image_1_1770.png
+image_1_1106.png
+image_1_1133.png
+image_1_1021.png
+image_0_1153.png
+image_1_1655.png
+image_0_1572.png
+image_0_1068.png
+image_0_1715.png
+image_1_1814.png
+image_0_1127.png
+image_1_1373.png
+image_1_1748.png
+image_1_1909.png
+image_0_1513.png
+image_1_1121.png
+image_0_1220.png
+image_1_1440.png
+image_0_1709.png
+image_1_1464.png
+image_0_1131.png
+image_0_1035.png
+image_0_1903.png
+image_1_1192.png
+image_1_1276.png
+image_0_1799.png
+image_0_1863.png
+image_0_1728.png
+image_1_1743.png
+image_0_1746.png
+image_0_1894.png
+image_1_1503.png
+image_0_1180.png
+image_1_1295.png
+image_0_1078.png
+image_1_1178.png
+image_1_1886.png
+image_0_1077.png
+image_1_1926.png
+image_0_1676.png
+image_0_1944.png
+image_0_1870.png
+image_0_1367.png
+image_0_1640.png
+image_1_1233.png
+image_1_1461.png
+image_0_1829.png
+image_1_1283.png
+image_0_1680.png
+image_0_1184.png
+image_0_1201.png
+image_0_1615.png
+image_0_1890.png
+image_1_1584.png
+image_0_1088.png
+image_0_1882.png
+image_1_1198.png
+image_1_1119.png
+image_1_1085.png
+image_0_1632.png
+image_0_1235.png
+image_0_1518.png
+image_0_1881.png
+image_1_1132.png
+image_1_1203.png
+image_0_1875.png
+image_0_1616.png
+image_1_1714.png
+image_1_1456.png
+image_0_1330.png
+image_1_1668.png
+image_0_1149.png
+image_1_1474.png
+image_1_1698.png
+image_1_1498.png
+image_1_1149.png
+image_1_1422.png
+image_0_1447.png
+image_0_1197.png
+image_0_1249.png
+image_0_1332.png
+image_1_1493.png
+image_1_1800.png
+image_1_1602.png
+image_1_1732.png
+image_0_1422.png
+image_1_1405.png
+image_0_1839.png
+image_1_1429.png
+image_1_1617.png
+image_1_1683.png
+image_0_1777.png
+image_0_1614.png
+image_0_1512.png
+image_1_1969.png
+image_0_1880.png
+image_1_1399.png
+image_1_1219.png
+image_1_1519.png
+image_0_1580.png
+image_1_1585.png
+image_0_1360.png
+image_1_1684.png
+image_1_1624.png
+image_1_1176.png
+image_0_1039.png
+image_0_1297.png
+image_0_1693.png
+image_0_1296.png
+image_1_1143.png

src/superpoint_pruning/distillation/lightning_trainer.py

@@ -0,0 +1,285 @@
+import argparse
+from pathlib import Path
+from typing import Any
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+import yaml
+import lightning as pl
+import os
+
+from superpoint_pruning.models.superpoint import SuperPoint
+from superpoint_pruning.distillation.utils import rescale_image, load_grayscale_image
+from superpoint_pruning.distillation.losses import (
+    detector_loss_simple,
+    descriptor_loss_simple,
+    detector_kd_kl,
+)
+from superpoint_pruning.paths import DEFAULT_CONFIG_PATH
+
+DATA_PATH_KEYS = (
+    "train_image_dir",
+    "train_image_ids_file",
+    "val_image_dir",
+    "val_image_ids_file",
+    "ground_truth_dir",
+)
+
+
+def resolve_path(base: Path, value: str | Path) -> Path:
+    path = Path(value)
+    if not path.is_absolute():
+        path = base / path
+    return path.resolve()
+
+
+def load_config(path: str | Path, data_root: Path | None = None) -> dict[str, Any]:
+    config_path = Path(path).resolve()
+    with open(config_path, encoding="utf-8") as f:
+        cfg = yaml.safe_load(f)
+
+    base = data_root.resolve() if data_root is not None else config_path.parent
+    for key in DATA_PATH_KEYS:
+        cfg["data"][key] = str(resolve_path(base, cfg["data"][key]))
+    cfg["trainer"]["default_root_dir"] = str(
+        resolve_path(config_path.parent, cfg["trainer"]["default_root_dir"])
+    )
+    return cfg
+
+
+class ImageFolderDataset(Dataset):
+    """Simple grayscale dataset from an image directory."""
+
+    def __init__(
+        self,
+        image_dir: str,
+        image_ids_file: str,
+        ground_truth_dir: str,
+        keypoints_file: str,
+        descriptors_file: str,
+        array_ids_file: str,
+        image_size: tuple[int, int],
+    ) -> None:
+        self.image_ids = Path(image_ids_file).read_text().splitlines()
+        self.image_dir = image_dir
+        self.ground_truth_dir = Path(ground_truth_dir)
+        self.keypoint_logits = np.load(
+            self.ground_truth_dir / keypoints_file, mmap_mode="r"
+        )
+        self.descriptor_logits = np.load(
+            self.ground_truth_dir / descriptors_file, mmap_mode="r"
+        )
+        self.image_size = image_size
+
+        array_ids_file = self.ground_truth_dir / array_ids_file
+        if array_ids_file.exists():
+            array_ids = array_ids_file.read_text().splitlines()
+            self.array_indices = {
+                image_id: idx for idx, image_id in enumerate(array_ids)
+            }
+        else:
+            if len(self.image_ids) > len(self.keypoint_logits):
+                raise ValueError(
+                    "The image id list is longer than the ground-truth arrays."
+                )
+            self.array_indices = {
+                image_id: idx for idx, image_id in enumerate(self.image_ids)
+            }
+
+    def __len__(self) -> int:
+        return len(self.image_ids)
+
+    def __getitem__(
+        self, idx: int
+    ) -> tuple[torch.Tensor, Any, torch.Tensor, torch.Tensor]:
+        image_id = self.image_ids[idx]
+        image_path = os.path.join(self.image_dir, image_id)
+        image = load_grayscale_image(image_path)
+        img, scale = rescale_image(image, self.image_size)
+        img = img[None, None].astype(np.float32)
+        img = torch.from_numpy(img)
+
+        array_idx = self.array_indices[image_id]
+        keypoint_logits = torch.from_numpy(
+            np.array(self.keypoint_logits[array_idx], copy=True)
+        )
+        descriptor_logits = torch.from_numpy(
+            np.array(self.descriptor_logits[array_idx], copy=True)
+        )
+        return img[0], scale, keypoint_logits, descriptor_logits
+
+
+class SuperPointDataModule(pl.LightningDataModule):
+    def __init__(self, cfg: dict[str, Any]) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.train_ds: ImageFolderDataset | None = None
+        self.val_ds: ImageFolderDataset | None = None
+
+    def setup(self, stage: str | None = None) -> None:
+        self.train_ds = ImageFolderDataset(
+            image_dir=self.cfg["train_image_dir"],
+            image_ids_file=self.cfg["train_image_ids_file"],
+            ground_truth_dir=self.cfg["ground_truth_dir"],
+            keypoints_file=self.cfg["keypoints_file"],
+            descriptors_file=self.cfg["descriptors_file"],
+            array_ids_file=self.cfg["array_ids_file"],
+            image_size=self.cfg["image_size"],
+        )
+        self.val_ds = ImageFolderDataset(
+            image_dir=self.cfg["val_image_dir"],
+            image_ids_file=self.cfg["val_image_ids_file"],
+            ground_truth_dir=self.cfg["ground_truth_dir"],
+            keypoints_file=self.cfg["keypoints_file"],
+            descriptors_file=self.cfg["descriptors_file"],
+            array_ids_file=self.cfg["array_ids_file"],
+            image_size=self.cfg["image_size"],
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        if self.train_ds is None:
+            raise RuntimeError("DataModule is not set up.")
+        return DataLoader(
+            self.train_ds,
+            batch_size=self.cfg["batch_size"],
+            shuffle=True,
+            num_workers=self.cfg["num_workers"],
+            pin_memory=self.cfg["pin_memory"],
+            drop_last=True,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        if self.val_ds is None:
+            raise RuntimeError("DataModule is not set up.")
+        return DataLoader(
+            self.val_ds,
+            batch_size=self.cfg["batch_size"],
+            shuffle=False,
+            num_workers=self.cfg["num_workers"],
+            pin_memory=self.cfg["pin_memory"],
+            drop_last=False,
+        )
+
+
+class SuperPointLightningModule(pl.LightningModule):
+
+    def __init__(self, cfg: dict[str, Any]) -> None:
+        super().__init__()
+        self.save_hyperparameters(cfg)
+        self.cfg = cfg
+        self.model = SuperPoint(
+            num_keypoints=cfg["model"]["num_keypoints"], return_dense=True
+        )
+        pruning_config = cfg["model"]["prune"]
+        self.model.prune_backbone(pruning_config)
+        print(self.model)
+
+    def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        keypoints, descriptors = self.model(image)
+        return keypoints, descriptors
+
+    def _compute_loss(
+        self,
+        keypoints: torch.Tensor,
+        descriptors: torch.Tensor,
+        keypoints_gt: torch.Tensor,
+        descriptors_gt: torch.Tensor,
+    ) -> torch.Tensor:
+        score_term_hard = detector_loss_simple(keypoints_gt, keypoints)
+        score_term = detector_kd_kl(keypoints_gt, keypoints)
+        desc_term = descriptor_loss_simple(descriptors, descriptors_gt)
+
+        self.log(
+            f"train/loss_ce",
+            score_term_hard,
+            prog_bar=True,
+            on_step=True,
+            on_epoch=True,
+        )
+        self.log(
+            f"train/loss_kl", score_term, prog_bar=True, on_step=True, on_epoch=True
+        )
+        self.log(
+            f"train/loss_desc", desc_term, prog_bar=True, on_step=True, on_epoch=True
+        )
+
+        return (
+            self.cfg["loss"]["cross_entropy_coef"] * score_term_hard
+            + self.cfg["loss"]["kl_coef"] * score_term
+            + self.cfg["loss"]["descriptor_coef"] * desc_term
+        )
+
+    def _shared_step(self, batch: dict[str, Any], stage: str) -> torch.Tensor:
+        images, scales, keypoints_gt, descriptors_gt = batch
+        keypoint_logits, descriptor_logits = self.forward(images)
+        loss = self._compute_loss(
+            keypoint_logits, descriptor_logits, keypoints_gt, descriptors_gt
+        )
+        self.log(
+            f"{stage}/loss",
+            loss,
+            prog_bar=True,
+            on_step=(stage == "train"),
+            on_epoch=True,
+        )
+        self.log(
+            f"{stage}/avg_score",
+            loss.mean(),
+            prog_bar=False,
+            on_step=False,
+            on_epoch=True,
+        )
+        return loss
+
+    def training_step(self, batch: dict[str, Any], batch_idx: int) -> torch.Tensor:
+        return self._shared_step(batch, stage="train")
+
+    def validation_step(self, batch: dict[str, Any], batch_idx: int) -> None:
+        self._shared_step(batch, stage="val")
+
+    def configure_optimizers(self) -> torch.optim.Optimizer:
+        trainable_params = [p for p in self.parameters() if p.requires_grad]
+        return torch.optim.AdamW(
+            trainable_params,
+            lr=self.cfg["optimizer"]["lr"],
+            weight_decay=self.cfg["optimizer"]["weight_decay"],
+        )
+
+
+def add_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--config", type=Path, default=DEFAULT_CONFIG_PATH, help="Path to YAML config"
+    )
+    parser.add_argument(
+        "--data-root",
+        type=Path,
+        default=None,
+        help="Base for relative data paths in the config. Default: directory of --config.",
+    )
+
+
+def main(args: argparse.Namespace) -> None:
+    cfg = load_config(args.config, data_root=args.data_root)
+
+    pl.seed_everything(cfg["seed"], workers=True)
+    data_module = SuperPointDataModule(cfg["data"])
+    lightning_module = SuperPointLightningModule(cfg)
+
+    trainer = pl.Trainer(
+        max_epochs=cfg["trainer"]["max_epochs"],
+        accelerator=cfg["trainer"]["accelerator"],
+        devices=cfg["trainer"]["devices"],
+        precision=cfg["trainer"]["precision"],
+        log_every_n_steps=cfg["trainer"]["log_every_n_steps"],
+        default_root_dir=cfg["trainer"]["default_root_dir"],
+        # enable_checkpointing=False,
+        logger=True,
+        limit_val_batches=0,
+    )
+    trainer.fit(model=lightning_module, datamodule=data_module)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    add_parser_args(parser)
+    main(parser.parse_args())

src/superpoint_pruning/distillation/losses.py

@@ -0,0 +1,159 @@
+import torch
+import torch.nn.functional as F
+
+
+def convert_keypoints_to_(
+    gt: torch.Tensor, H: int = 768, W: int = 1024
+) -> torch.Tensor:
+    B, N, _ = gt.shape
+    xy = gt.round().long()
+    x = xy[..., 0].clamp(0, W - 1)
+    y = xy[..., 1].clamp(0, H - 1)
+    keypoint_map = torch.zeros((B, H, W), device=gt.device, dtype=torch.float32)
+    b = torch.arange(B, device=gt.device).unsqueeze(1).expand(B, N)
+    keypoint_map[b, y, x] = 1.0
+    return keypoint_map
+
+
+def detector_loss(
+    keypoint_map: torch.Tensor,  # (B,H,W) or (B,1,H,W), binary/bool
+    logits: torch.Tensor,  # (B,65,Hc,Wc), raw convPb output
+    valid_mask: torch.Tensor | None = None,  # (B,H,W) or (B,1,H,W)
+    grid_size: int = 8,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    # --- labels: space_to_depth + dustbin + random tie-break + argmax ---
+    if keypoint_map.ndim == 3:
+        keypoint_map = keypoint_map.unsqueeze(1)  # (B,1,H,W)
+    keypoint_map = keypoint_map.float()
+    # TF NHWC space_to_depth -> PyTorch NCHW pixel_unshuffle
+    labels = F.pixel_unshuffle(keypoint_map, downscale_factor=grid_size)  # (B,64,Hc,Wc)
+    dustbin = torch.ones_like(labels[:, :1])  # (B,1,Hc,Wc)
+    labels = torch.cat([2.0 * labels, dustbin], dim=1)  # (B,65,Hc,Wc)
+    # same tie-break idea as TF random_uniform(..., 0, 0.1)
+    labels = torch.argmax(
+        labels + 0.1 * torch.rand_like(labels), dim=1
+    )  # (B,Hc,Wc), long
+    # --- valid mask path ---
+    if valid_mask is None:
+        valid_mask = torch.ones_like(keypoint_map)
+    elif valid_mask.ndim == 3:
+        valid_mask = valid_mask.unsqueeze(1)
+    valid_mask = valid_mask.float()
+    valid_mask = F.pixel_unshuffle(
+        valid_mask, downscale_factor=grid_size
+    )  # (B,64,Hc,Wc)
+    valid_mask = torch.prod(valid_mask, dim=1)  # (B,Hc,Wc)
+    # --- sparse softmax cross entropy with weights ---
+    per_cell = F.cross_entropy(logits, labels, reduction="none")  # (B,Hc,Wc)
+    weighted = per_cell * valid_mask
+    loss = weighted.sum() / valid_mask.sum().clamp_min(1.0 + eps)
+    return loss
+
+
+def detector_loss_simple(
+    gt_logits: torch.Tensor,
+    pred_logits: torch.Tensor,  # (B,65,Hc,Wc), raw convPb output
+) -> torch.Tensor:
+    gt_labels = torch.argmax(
+        gt_logits + 0.1 * torch.rand_like(gt_logits), dim=1
+    )  # (B,Hc,Wc), long
+    return F.cross_entropy(pred_logits, gt_labels, reduction="mean")
+
+
+def detector_kd_kl(teacher_logits, student_logits, T=2.0, valid_mask=None, eps=1e-8):
+    # student_logits, teacher_logits: (B,65,Hc,Wc)
+    log_p_s = F.log_softmax(student_logits / T, dim=1)
+    p_t = F.softmax(teacher_logits / T, dim=1)
+
+    # KL per cell: (B,Hc,Wc)
+    kl_map = F.kl_div(log_p_s, p_t, reduction="none").sum(dim=1)
+
+    if valid_mask is None:
+        return (T * T) * kl_map.mean()
+
+    # valid_mask expected (B,H,W) or (B,1,H,W), convert to cell mask (B,Hc,Wc)
+    if valid_mask.ndim == 3:
+        valid_mask = valid_mask.unsqueeze(1)
+    vm = F.pixel_unshuffle(valid_mask.float(), downscale_factor=8)  # (B,64,Hc,Wc)
+    vm = torch.prod(vm, dim=1)  # (B,Hc,Wc)
+
+    return (T * T) * (kl_map * vm).sum() / vm.sum().clamp_min(1.0 + eps)
+
+
+def descriptor_loss_simple(
+    pred_desc: torch.Tensor,  # (B, D, Hc, Wc)
+    gt_desc: torch.Tensor,  # (B, D, Hc, Wc)
+    valid_mask: torch.Tensor | None = None,  # (B,H,W) or (B,1,H,W)
+    grid_size: int = 8,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    pred = F.normalize(pred_desc, p=2, dim=1)
+    gt = F.normalize(gt_desc, p=2, dim=1)
+
+    # cosine distance per cell
+    per_cell = 1.0 - (pred * gt).sum(dim=1)  # (B, Hc, Wc)
+
+    if valid_mask is None:
+        return per_cell.mean()
+
+    if valid_mask.ndim == 3:
+        valid_mask = valid_mask.unsqueeze(1)  # (B,1,H,W)
+
+    vm = F.pixel_unshuffle(
+        valid_mask.float(), downscale_factor=grid_size
+    )  # (B,64,Hc,Wc)
+    vm = torch.prod(vm, dim=1)  # (B,Hc,Wc)
+
+    return (per_cell * vm).sum() / vm.sum().clamp_min(1.0 + eps)
+
+
+def descriptor_loss(
+    descriptors: torch.Tensor,  # (B, D, Hc, Wc), student
+    target_descriptors: torch.Tensor,  # (B, D, Hc, Wc), teacher/GT
+    valid_mask: torch.Tensor | None = None,  # (B,H,W) or (B,1,H,W)
+    grid_size: int = 8,
+    positive_margin: float = 1.0,
+    negative_margin: float = 0.2,
+    lambda_d: float = 0.05,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    B, D, Hc, Wc = descriptors.shape
+    HW = Hc * Wc
+    # L2 normalize descriptors
+    desc = F.normalize(descriptors, p=2, dim=1)  # (B,D,Hc,Wc)
+    tgt = F.normalize(target_descriptors, p=2, dim=1)  # (B,D,Hc,Wc)
+    # Flatten spatial dims
+    desc = desc.flatten(2).transpose(1, 2)  # (B,HW,D)
+    tgt = tgt.flatten(2).transpose(1, 2)  # (B,HW,D)
+    # Pairwise dot products: (B,HW,HW)
+    dot = torch.bmm(desc, tgt.transpose(1, 2))
+    dot = F.relu(dot)
+    # TF does double normalization over pairwise axes
+    dot = F.normalize(dot, p=2, dim=2)
+    dot = F.normalize(dot, p=2, dim=1)
+    # Identity correspondence mask s (diagonal)
+    eye = torch.eye(HW, device=dot.device, dtype=dot.dtype).unsqueeze(0)  # (1,HW,HW)
+    s = eye.expand(B, -1, -1)
+    positive_dist = F.relu(positive_margin - dot)
+    negative_dist = F.relu(dot - negative_margin)
+    pairwise_loss = (
+        lambda_d * s * positive_dist + (1.0 - s) * negative_dist
+    )  # (B,HW,HW)
+    # valid mask: same logic as TF space_to_depth + reduce_prod
+    if valid_mask is None:
+        vm = torch.ones(
+            (B, Hc * grid_size, Wc * grid_size), device=dot.device, dtype=dot.dtype
+        )
+    else:
+        vm = valid_mask
+        if vm.ndim == 3:
+            vm = vm.unsqueeze(1)  # (B,1,H,W)
+        vm = vm.float()
+        vm = F.pixel_unshuffle(vm, downscale_factor=grid_size)  # (B,grid^2,Hc,Wc)
+        vm = torch.prod(vm, dim=1)  # (B,Hc,Wc)
+    vm = vm.reshape(B, HW)  # valid target cells
+    vm = vm[:, None, :]  # (B,1,HW), broadcast to (B,HW,HW)
+    normalization = vm.sum() * float(HW) + eps
+    loss = (pairwise_loss * vm).sum() / normalization
+    return loss

src/superpoint_pruning/distillation/setup.py

@@ -0,0 +1,161 @@
+import random
+import os
+from pathlib import Path
+from superpoint_pruning.distillation.utils import rescale_image, load_grayscale_image
+from superpoint_pruning.models.superpoint import SuperPoint
+from superpoint_pruning.paths import DATA_ROOT, DEFAULT_DATASET_NAME
+import numpy as np
+import torch
+from tqdm import tqdm
+
+import urllib.request
+import zipfile
+import argparse
+
+DEFAULT_DATASET_URL = "http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/indoor_forward_3_snapdragon_with_gt.zip"
+DEFAULT_GT_INDEX_FILE = "train_base.txt"
+
+
+def download_dataset_from_url(url: str, output_dir: Path) -> None:
+
+    zip_path = DATA_ROOT / "download.zip"
+    if output_dir.exists():
+        print(
+            f"Dataset directory {output_dir.resolve()} already exists. Skipping download."
+        )
+        return
+    output_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading dataset from {url} to {zip_path.resolve()}")
+    urllib.request.urlretrieve(url, zip_path)
+    print(f"Extracting files from {zip_path.resolve()} to {output_dir.resolve()}")
+    with zipfile.ZipFile(zip_path, "r") as archive:
+        archive.extractall(output_dir)
+    zip_path.unlink()
+    print(f"Files extracted to {output_dir.resolve()}")
+
+
+def generate_data_split(
+    image_dir: Path, split_ratio: float = 0.8
+) -> tuple[list[str], list[str]]:
+    random.seed(42)
+    images = [path.name for path in image_dir.glob("*.png")]
+    random.shuffle(images)
+    split_index = int(len(images) * split_ratio)
+    train_images = images[:split_index]
+    val_images = images[split_index:]
+
+    with open(DATA_ROOT / "train.txt", "w") as f:
+        for image in train_images:
+            f.write(image + "\n")
+    with open(DATA_ROOT / "val.txt", "w") as f:
+        for image in val_images:
+            f.write(image + "\n")
+
+    print(
+        f"Train images: {len(train_images)}, Val images: {len(val_images)} - Saved to {DATA_ROOT}"
+    )
+
+
+@torch.inference_mode()
+def generate_ground_truth(
+    image_dir: Path,
+    id_file: Path,
+    num_keypoints: int = 512,
+    image_size: tuple[int, int] = (640, 480),
+    num_images: int = 250,
+) -> None:
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    images = [path.name for path in image_dir.glob("*.png")]
+    model = SuperPoint(num_keypoints=num_keypoints, return_dense=True)
+    model.eval().to(device)
+
+    ids = id_file.read_text().splitlines()
+    ids = ids[:num_images]
+    id_set = set(ids)
+    selected_images = [image_id for image_id in images if image_id in id_set]
+    Path(DATA_ROOT / f"gt_id_file_{Path(id_file).stem}_{num_images}.txt").write_text(
+        "\n".join(selected_images) + "\n"
+    )
+
+    desc = None
+    kpts = None
+
+    for index, image_id in enumerate(tqdm(selected_images)):
+        image_path = os.path.join(image_dir, image_id)
+        image = load_grayscale_image(image_path)
+        img, _ = rescale_image(image, image_size)
+        img = img[None, None].astype(np.float32)
+
+        img = torch.from_numpy(img).to(device)
+
+        keypoint_logits, descriptor_logits = model(img)
+
+        desc_np = descriptor_logits.detach().cpu().numpy()[0]
+        kpts_np = keypoint_logits.detach().cpu().numpy()[0]
+
+        if desc is None or kpts is None:
+            desc = np.lib.format.open_memmap(
+                DATA_ROOT / f"desc_{id_file.stem}_{num_images}.npy",
+                mode="w+",
+                dtype=desc_np.dtype,
+                shape=(len(selected_images), *desc_np.shape),
+            )
+            kpts = np.lib.format.open_memmap(
+                DATA_ROOT / f"kpts_{id_file.stem}_{num_images}.npy",
+                mode="w+",
+                dtype=kpts_np.dtype,
+                shape=(len(selected_images), *kpts_np.shape),
+            )
+
+        desc[index] = desc_np
+        kpts[index] = kpts_np
+
+    if desc is not None:
+        desc.flush()
+    if kpts is not None:
+        kpts.flush()
+
+    print(f"Ground truth files generated to {DATA_ROOT}")
+
+
+def add_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--download-url", type=str, default=DEFAULT_DATASET_URL)
+    parser.add_argument("--output-dir-name", type=str, default=DEFAULT_DATASET_NAME)
+    parser.add_argument("--gt-index-file", type=str, default=DEFAULT_GT_INDEX_FILE)
+    parser.add_argument("--generate-data-split", action="store_true", default=False)
+    parser.add_argument("--generate-gt", action="store_true", default=False)
+
+    parser.add_argument("--num-keypoints", type=int, default=512)
+    parser.add_argument("--width", type=int, default=640)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--num-images", type=int, default=250)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Setup the dataset for distillation and evaluation."
+    )
+    add_parser_args(parser)
+    return parser
+
+
+def main(args: argparse.Namespace) -> None:
+    if args.download_url and args.output_dir_name:
+        download_dataset_from_url(
+            args.download_url, DATA_ROOT / "datasets" / args.output_dir_name
+        )
+    if args.generate_data_split:
+        generate_data_split(DATA_ROOT / "datasets" / args.output_dir_name / "img", 0.8)
+    if args.generate_gt:
+        generate_ground_truth(
+            DATA_ROOT / "datasets" / args.output_dir_name / "img",
+            DATA_ROOT / args.gt_index_file,
+            num_images=args.num_images,
+            image_size=(args.width, args.height),
+            num_keypoints=args.num_keypoints,
+        )
+
+
+if __name__ == "__main__":
+    main(build_parser().parse_args())

src/superpoint_pruning/distillation/utils.py

@@ -0,0 +1,74 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import cv2
+
+
+def plot_images(images, titles=None, figsize=None):
+    n = len(images)
+    if figsize is None:
+        figsize = (5 * n, 5)
+
+    fig, axes = plt.subplots(1, n, figsize=figsize)
+    axes = np.atleast_1d(axes)
+
+    for i, (ax, image) in enumerate(zip(axes, images)):
+        if image.ndim == 2:
+            ax.imshow(image, cmap="gray")
+        else:
+            ax.imshow(image)
+        if titles is not None:
+            ax.set_title(titles[i])
+        ax.axis("off")
+
+    fig.tight_layout()
+    return fig, axes
+
+
+def plot_images_with_keypoints(
+    images,
+    keypoints,
+    titles=None,
+    keypoint_size=30,
+    figsize=None,
+):
+    fig, axes = plot_images(
+        images,
+        titles=titles,
+        figsize=figsize,
+    )
+    for ax, points in zip(axes, keypoints):
+        points = np.asarray(points)
+        if len(points) == 0:
+            continue
+        x = points[:, 0]
+        y = points[:, 1]
+
+        ax.scatter(x, y, s=keypoint_size, marker=".", color="lime")
+    fig.tight_layout()
+
+
+def load_grayscale_image(path):
+    image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+
+    if image is None:
+        raise FileNotFoundError(f"Could not read image: {path}")
+
+    image = image.astype(np.float32) / 255.0
+
+    return image
+
+
+def rescale_image(image, new_size=None):
+    height, width = image.shape[:2]
+    new_width, new_height = new_size
+
+    resized = cv2.resize(
+        image,
+        (new_width, new_height),
+        interpolation=cv2.INTER_LINEAR,
+    )
+
+    scale_x = new_width / width
+    scale_y = new_height / height
+
+    return resized, (scale_x, scale_y)

src/superpoint_pruning/evaluation/benchmark_sp_tensorrt.py

@@ -0,0 +1,178 @@
+import argparse
+import numpy as np
+from tqdm import tqdm
+import time
+import os
+import threading
+from queue import Queue, Empty
+from pathlib import Path
+
+from superpoint_pruning.distillation.utils import load_grayscale_image, rescale_image
+from superpoint_pruning.paths import DEFAULT_IMAGE_DIR
+
+QUEUE_SENTINEL = object()
+
+
+def benchmark_tensorrt(
+    img_dir: str,
+    model_path: str,
+    max_keypoints: int = 512,
+    descriptor_dim: int = 256,
+    num_images: int = 1000,
+    start_index: int = 1000,
+    num_loader_threads: int = 2,
+    queue_size: int = 32,
+    image_size: tuple = (640, 480),
+):
+
+    try:
+        from superpoint_pruning.evaluation.onnx_helper import SP_ONNXClassifierWrapper
+    except ImportError:
+        raise ImportError(
+            "ONNX helper requires pycuda and TensorRT utilities installed."
+        )
+
+    trt_model = SP_ONNXClassifierWrapper(
+        str(model_path), max_keypoints=max_keypoints, descriptor_dim=descriptor_dim
+    )
+    times = []
+    index_queue = Queue()
+    data_queue = Queue(maxsize=max(1, queue_size))
+    errors = []
+    errors_lock = threading.Lock()
+    stop_event = threading.Event()
+    processed_images = 0
+
+    for img_index in range(start_index, start_index + num_images):
+        index_queue.put(img_index)
+
+    effective_loader_threads = max(1, num_loader_threads)
+
+    def loader_worker():
+        try:
+            while not stop_event.is_set():
+                try:
+                    img_index = index_queue.get_nowait()
+                except Empty:
+                    break
+
+                img_path = os.path.join(img_dir, f"image_0_{img_index}.png")
+                original = load_grayscale_image(img_path)
+                img, scale = rescale_image(original, new_size=image_size)
+                img = img[None, None].astype(np.float32)
+                data_queue.put((img, scale))
+        except Exception as exc:
+            with errors_lock:
+                errors.append(exc)
+            stop_event.set()
+        finally:
+            data_queue.put(QUEUE_SENTINEL)
+
+    loader_threads = [
+        threading.Thread(target=loader_worker, daemon=True)
+        for _ in range(effective_loader_threads)
+    ]
+
+    for thread in loader_threads:
+        thread.start()
+    try:
+        finished_loaders = 0
+        with tqdm(total=num_images, desc="Inferencing", unit="img") as pbar:
+            while finished_loaders < effective_loader_threads:
+                if stop_event.is_set() and data_queue.empty():
+                    break
+
+                try:
+                    item = data_queue.get(timeout=0.1)
+                except Empty:
+                    continue
+
+                if item is QUEUE_SENTINEL:
+                    finished_loaders += 1
+                    continue
+
+                img, scale = item
+                scale = np.array(scale, dtype=np.float32)
+                start = time.perf_counter()
+                keypoints, _, descriptors = trt_model.predict(img)
+                keypoints = (keypoints.astype(np.float32) + 0.5) / scale[None] - 0.5
+                stop = time.perf_counter()
+                times.append(stop - start)
+                processed_images += 1
+                pbar.update(1)
+
+        for thread in loader_threads:
+            thread.join()
+
+        if errors:
+            raise RuntimeError(f"Benchmark failed in loader thread: {errors[0]}")
+
+        if processed_images != num_images:
+            raise RuntimeError(
+                f"Processed {processed_images}/{num_images} images before stopping."
+            )
+
+        print(f"Average time: {np.array(times).mean()}")
+    finally:
+        trt_model.close()
+
+
+def add_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--img-dir", type=Path, default=DEFAULT_IMAGE_DIR)
+    parser.add_argument(
+        "--model-path",
+        type=Path,
+        required=True,
+        help="Path to the TensorRT SuperPoint engine.",
+    )
+    parser.add_argument(
+        "--max-keypoints",
+        type=int,
+        default=512,
+        help="Maximum number of keypoints from input.",
+    )
+    parser.add_argument(
+        "--descriptor-dim", type=int, default=256, help="Descriptor dimension."
+    )
+    parser.add_argument(
+        "--num-images", type=int, default=1000, help="Number of images to benchmark."
+    )
+    parser.add_argument(
+        "--start-index", type=int, default=1000, help="Start index of the images."
+    )
+    parser.add_argument(
+        "--num-loader-threads",
+        type=int,
+        default=2,
+        help="Number of producer threads for loading and preprocessing images.",
+    )
+    parser.add_argument(
+        "--queue-size",
+        type=int,
+        default=64,
+        help="Max number of preprocessed images buffered for inference.",
+    )
+    parser.add_argument("--width", type=int, default=640, help="Image width.")
+    parser.add_argument("--height", type=int, default=480, help="Image height.")
+
+
+def main(args: argparse.Namespace) -> None:
+    benchmark_tensorrt(
+        img_dir=args.img_dir,
+        model_path=args.model_path,
+        max_keypoints=args.max_keypoints,
+        descriptor_dim=args.descriptor_dim,
+        num_images=args.num_images,
+        start_index=args.start_index,
+        num_loader_threads=args.num_loader_threads,
+        queue_size=args.queue_size,
+        image_size=(args.width, args.height),
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Benchmark a TensorRT SuperPoint engine."
+    )
+    add_parser_args(parser)
+    main(parser.parse_args())

src/superpoint_pruning/evaluation/eval.py

@@ -0,0 +1,293 @@
+import argparse
+from pathlib import Path
+
+import torch
+import numpy as np
+from tqdm import tqdm
+
+from superpoint_pruning.models.superpoint import SuperPoint
+from superpoint_pruning.distillation.utils import (
+    load_grayscale_image,
+    rescale_image,
+    plot_images_with_keypoints,
+)
+from superpoint_pruning.evaluation.metrics import BenchmarkMetrics
+from superpoint_pruning.paths import DEFAULT_IMAGE_DIR, DEFAULT_TRAINING_IDS
+from superpoint_pruning.pruning import add_pruning_parser_args, pruning_config_from_args
+
+
+def _image_path(image_dir: Path | str, sample: int) -> Path:
+    return Path(image_dir) / f"image_0_{sample}.png"
+
+
+@torch.inference_mode()
+def evaluate(
+    sp_original,
+    sp_pruned,
+    sp_head,
+    matcher,
+    metrics,
+    image_dir,
+    start_idx=1000,
+    end_idx=2000,
+    device="cpu",
+    save_preds=False,
+    trained=False,
+    skip=False,
+    image_size=(1024, 768),
+    training_ids=None,
+    hierarchical=False,
+    skip_refinement=False,
+):
+
+    try:
+        from lightglue import LightGlue
+    except ImportError as exc:
+        raise ImportError(
+            "The evaluation requires LightGlue for matching metrics"
+        ) from exc
+
+    if training_ids is not None:
+        with open(training_ids, "r") as f:
+            image_ids = f.read().splitlines()
+            image_ids = [int(id.split("_")[-1].split(".")[0]) for id in image_ids]
+
+    skipped = 0
+    sp_head_hier = sp_pruned.dense_head
+    image_dir = Path(image_dir)
+
+    for sample in tqdm(range(start_idx, end_idx)):
+        if skip and sample in image_ids:
+            skipped += 1
+            continue
+        img = load_grayscale_image(str(_image_path(image_dir, sample)))
+        img, scale = rescale_image(img, new_size=image_size)
+        scale = torch.as_tensor(scale, device=device, dtype=torch.float32)
+        img = img[None, None].astype(np.float32)
+        img = torch.from_numpy(img).to(device)
+
+        kpts_dense, desc_dense = sp_pruned(img)
+        original_kpts_dense, original_desc_dense = sp_original(img)
+        kpts, _, desc = sp_head_hier(original_kpts_dense, desc_dense)
+        kpts = (kpts.to(torch.float32) + 0.5) / scale - 0.5
+
+        img2 = load_grayscale_image(str(_image_path(image_dir, sample + 10)))
+        img2, scale = rescale_image(img2, new_size=image_size)
+        scale = torch.as_tensor(scale, device=device, dtype=torch.float32)
+        img2 = img2[None, None].astype(np.float32)
+        img2 = torch.from_numpy(img2).to(device)
+
+        kpts_dense2, desc_dense2 = sp_pruned(img2)
+        original_kpts_dense2, original_desc_dense2 = sp_original(img2)
+        kpts2, _, desc2 = sp_head_hier(original_kpts_dense2, desc_dense2)
+        kpts2 = (kpts2.to(torch.float32) + 0.5) / scale - 0.5
+
+        original_kpts, _, original_desc = sp_head(
+            original_kpts_dense, original_desc_dense
+        )
+        original_kpts = (original_kpts.to(torch.float32) + 0.5) / scale - 0.5
+        original_kpts2, _, original_desc2 = sp_head(
+            original_kpts_dense2, original_desc_dense2
+        )
+        original_kpts2 = (original_kpts2.to(torch.float32) + 0.5) / scale - 0.5
+        pruned_kpts, _, pruned_desc = sp_head_hier(kpts_dense, desc_dense)
+        pruned_kpts = (pruned_kpts.to(torch.float32) + 0.5) / scale - 0.5
+
+        pruned_kpts2, _, pruned_desc2 = sp_head_hier(kpts_dense2, desc_dense2)
+        pruned_kpts2 = (pruned_kpts2.to(torch.float32) + 0.5) / scale - 0.5
+        metrics.update_keypoints(
+            pruned_kpts.detach().cpu(), original_kpts.detach().cpu()
+        )
+
+        original_matches = matcher(
+            {
+                "image0": {
+                    "keypoints": original_kpts,
+                    "descriptors": original_desc,
+                    "image_size": torch.tensor([[640.0, 480.0]], device=device),
+                },
+                "image1": {
+                    "keypoints": original_kpts2,
+                    "descriptors": original_desc2,
+                    "image_size": torch.tensor([[640.0, 480.0]], device=device),
+                },
+            }
+        )["matches"][0]
+
+        matches = matcher(
+            {
+                "image0": {
+                    "keypoints": kpts,
+                    "descriptors": desc,
+                    "image_size": torch.tensor([[640.0, 480.0]], device=device),
+                },
+                "image1": {
+                    "keypoints": kpts2,
+                    "descriptors": desc2,
+                    "image_size": torch.tensor([[640.0, 480.0]], device=device),
+                },
+            }
+        )["matches"][0]
+
+        coord_matches = torch.cat(
+            (kpts[0][matches[:, 0]], kpts2[0][matches[:, 1]]), dim=1
+        )
+        coord_original_matches = torch.cat(
+            (
+                original_kpts[0][original_matches[:, 0]],
+                original_kpts2[0][original_matches[:, 1]],
+            ),
+            dim=1,
+        )
+        matches_pruned = matcher(
+            {
+                "image0": {
+                    "keypoints": pruned_kpts,
+                    "descriptors": pruned_desc,
+                    "image_size": torch.tensor([[640.0, 480.0]], device=device),
+                },
+                "image1": {
+                    "keypoints": pruned_kpts2,
+                    "descriptors": pruned_desc2,
+                    "image_size": torch.tensor([[640.0, 480.0]], device=device),
+                },
+            }
+        )["matches"][0]
+        metrics.update_matches(
+            coord_matches.detach().cpu(), coord_original_matches.detach().cpu()
+        )
+        if hasattr(metrics, "pruned_matches"):
+            metrics.pruned_matches.append(len(matches_pruned.detach().cpu()))
+        else:
+            metrics.pruned_matches = [len(matches_pruned.detach().cpu())]
+
+    metrics.print_metrics()
+    print(
+        f"Average number of pruned matches (pruned keypoints + pruned descriptors): {np.mean(np.array(metrics.pruned_matches))}"
+    )
+    print("\n")
+    print(f"Skipped {skipped} images")
+
+
+@torch.inference_mode()
+def evaluate_simple(
+    sp_original,
+    sp_pruned,
+    sp_head,
+    metrics,
+    image_dir,
+    start_idx=1000,
+    end_idx=2000,
+    device="cpu",
+    save_preds=False,
+    trained=False,
+    skip=False,
+    image_size=(1024, 768),
+    training_ids=None,
+    hierarchical=False,
+    skip_refinement=False,
+):
+
+    if training_ids is not None:
+        with open(training_ids, "r") as f:
+            image_ids = f.read().splitlines()
+            image_ids = [int(id.split("_")[-1].split(".")[0]) for id in image_ids]
+
+    skipped = 0
+    desc_l2s = []
+    sp_head_hier = sp_pruned.dense_head
+    image_dir = Path(image_dir)
+
+    for sample in tqdm(range(start_idx, end_idx)):
+        if skip and sample in image_ids:
+            skipped += 1
+            continue
+        img = load_grayscale_image(str(_image_path(image_dir, sample)))
+        img, scale = rescale_image(img, new_size=image_size)
+        scale = torch.as_tensor(scale, device=device, dtype=torch.float32)
+        img = img[None, None].astype(np.float32)
+        img = torch.from_numpy(img).to(device)
+
+        kpts_dense, desc_dense = sp_pruned(img)
+        original_kpts_dense, original_desc_dense = sp_original(img)
+        kpts, _, desc = sp_head_hier(original_kpts_dense, desc_dense)
+        kpts = (kpts.to(torch.float32) + 0.5) / scale - 0.5
+
+        original_kpts, _, original_desc = sp_head(
+            original_kpts_dense, original_desc_dense
+        )
+        original_kpts = (original_kpts.to(torch.float32) + 0.5) / scale - 0.5
+        pruned_kpts, _, pruned_desc = sp_head_hier(kpts_dense, desc_dense)
+        pruned_kpts = (pruned_kpts.to(torch.float32) + 0.5) / scale - 0.5
+
+        desc_l2 = torch.norm(desc - original_desc, p=2, dim=-1).mean()
+        desc_l2s.append(desc_l2.item())
+
+        metrics.update_keypoints(
+            pruned_kpts.detach().cpu(), original_kpts.detach().cpu()
+        )
+
+    metrics.print_metrics()
+    print(f"Average descriptor L2: {np.mean(desc_l2s):.6f}")
+    print(f"Skipped {skipped} images")
+
+
+def add_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--image-dir", type=Path, default=DEFAULT_IMAGE_DIR)
+    parser.add_argument("--training-ids", type=Path, default=DEFAULT_TRAINING_IDS)
+    parser.add_argument("--start-idx", type=int, default=1000)
+    parser.add_argument("--end-idx", type=int, default=2000)
+    parser.add_argument("--num-keypoints", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=640)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--skip", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument(
+        "--hierarchical", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--skip-refinement", action=argparse.BooleanOptionalAction, default=False
+    )
+    add_pruning_parser_args(parser)
+
+
+def main(args: argparse.Namespace) -> None:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pruning_config, pruning_checkpoint = pruning_config_from_args(args)
+    sp_original = (
+        SuperPoint(num_keypoints=args.num_keypoints, return_dense=True)
+        .eval()
+        .to(device)
+    )
+    sp_pruned = SuperPoint(
+        num_keypoints=args.num_keypoints,
+        return_dense=True,
+        hierarchical_topk=args.hierarchical,
+        skip_refinement=args.skip_refinement,
+    ).eval()
+    if pruning_config:
+        sp_pruned.prune_backbone(pruning_config)
+    if pruning_checkpoint is not None:
+        sp_pruned.load_pruned_weights(str(pruning_checkpoint))
+    sp_pruned = sp_pruned.to(device).eval()
+
+    evaluate_simple(
+        sp_original,
+        sp_pruned,
+        sp_original.dense_head,
+        BenchmarkMetrics(),
+        args.image_dir,
+        start_idx=args.start_idx,
+        end_idx=args.end_idx,
+        device=device,
+        skip=args.skip,
+        training_ids=args.training_ids,
+        image_size=(args.width, args.height),
+        hierarchical=args.hierarchical,
+        skip_refinement=args.skip_refinement,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Evaluate a SuperPoint model.")
+    add_parser_args(parser)
+    main(parser.parse_args())

src/superpoint_pruning/evaluation/metrics.py

@@ -0,0 +1,106 @@
+import numpy as np
+import torch
+
+
+def match_overlap(pred, gt):
+    pred = pred.detach().cpu() if torch.is_tensor(pred) else torch.as_tensor(pred)
+    gt = gt.detach().cpu() if torch.is_tensor(gt) else torch.as_tensor(gt)
+
+    if len(pred.shape) > 1 and pred.shape[1] == 3:
+        pred = pred[:, 1:]
+    if len(gt.shape) > 1 and gt.shape[1] == 3:
+        gt = gt[:, 1:]
+
+    if len(pred) > 0:
+        invalid = torch.nonzero(pred[:, 0] == -1, as_tuple=False)
+        if len(invalid) > 0:
+            pred = pred[: invalid[0].item()]
+
+    num_preds = len(pred)
+    if num_preds == 0 or len(gt) == 0:
+        running_recall = 0
+    else:
+        overlap = (pred[:, None, :] == gt[None, :, :]).all(dim=2)
+        running_recall = overlap.any(dim=1).sum().item()
+
+    recall = running_recall / len(gt) if len(gt) > 0 else 0
+    precision = running_recall / num_preds if num_preds > 0 else 0
+
+    # recall: # of original matches found
+    # precision: in this case, the # of correct predictions
+    return recall, precision, num_preds
+
+
+def keypoint_overlap(pred, gt, cell_size=8, max_x=640, max_y=480):
+    if len(pred.shape) == 3:
+        pred = pred[0]
+    if len(gt.shape) == 3:
+        gt = gt[0]
+
+    num_cells_x = (max_x + cell_size - 1) // cell_size
+    num_cells_y = (max_y + cell_size - 1) // cell_size
+    num_cells = num_cells_x * num_cells_y
+
+    def counts_per_cell(kpts):
+        valid = (
+            (kpts[:, 0] >= 0)
+            & (kpts[:, 0] < max_x)
+            & (kpts[:, 1] >= 0)
+            & (kpts[:, 1] < max_y)
+        )
+        kpts = kpts[valid]
+        cell_x = torch.div(kpts[:, 0], cell_size, rounding_mode="floor").long()
+        cell_y = torch.div(kpts[:, 1], cell_size, rounding_mode="floor").long()
+        linear_idx = cell_y * num_cells_x + cell_x
+        return torch.bincount(linear_idx, minlength=num_cells)
+
+    pred_counts = counts_per_cell(pred)
+    gt_counts = counts_per_cell(gt)
+    covered = torch.minimum(pred_counts, gt_counts).sum().item()
+    return covered, pred_counts.sum().item(), gt_counts.sum().item()
+
+
+class BenchmarkMetrics:
+    def __init__(self):
+        self.recalls = []
+        self.precisions = []
+        self.gt_counts = []
+        self.pred_counts = []
+        self.keypoints_covered = []
+        self.pred_counts_kpts = []
+        self.gt_counts_kpts = []
+
+    def update_matches(self, pred, gt):
+        recall, precision, num_preds = match_overlap(pred, gt)
+        self.recalls.append(recall)
+        self.precisions.append(precision)
+        self.gt_counts.append(len(gt))
+        self.pred_counts.append(num_preds)
+
+    def update_keypoints(self, pred, gt):
+        covered, pred_counts, gt_counts = keypoint_overlap(pred, gt)
+        self.keypoints_covered.append(covered)
+        self.pred_counts_kpts.append(pred_counts)
+        self.gt_counts_kpts.append(gt_counts)
+
+    def print_metrics(self):
+        print(f"Recall: {np.mean(self.recalls)}")
+        print(f"Precision: {np.mean(self.precisions)}")
+        print(
+            f"Average number of matches (original keypoints + pruned descriptors): {np.mean(np.array(self.pred_counts))}"
+        )
+        print(
+            f"Average number of matches (original keypoints + original descriptors): {np.mean(np.array(self.gt_counts))}"
+        )
+        print(
+            f"Average change in number of matches: {np.mean(np.array(self.pred_counts) - np.array(self.gt_counts))}"
+        )
+        print(
+            f"Average number of keypoints covered: {np.mean(np.array(self.keypoints_covered))}"
+        )
+        print(
+            f"Average number of keypoints in prediction: {np.mean(np.array(self.pred_counts_kpts))}"
+        )
+        print(
+            f"Average number of keypoints in ground truth: {np.mean(np.array(self.gt_counts_kpts))}"
+        )

src/superpoint_pruning/evaluation/onnx_helper.py

@@ -0,0 +1,97 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 numpy as np
+import tensorrt as trt
+
+import pycuda.driver as cuda
+import pycuda.autoinit
+
+
+class WrapperBase:
+
+    def close(self):
+        pass
+
+
+class SP_ONNXClassifierWrapper(WrapperBase):
+    def __init__(
+        self,
+        file,
+        target_dtype=np.float32,
+        batch_size=1,
+        max_keypoints=512,
+        descriptor_dim=256,
+    ):
+
+        self.target_dtype = target_dtype
+        self.load(file)
+
+        self.stream = None
+        self.max_keypoints = max_keypoints
+        self.batch_size = batch_size
+        self.descriptor_dim = descriptor_dim
+
+    def load(self, file):
+        f = open(file, "rb")
+        self.runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
+
+        self.engine = self.runtime.deserialize_cuda_engine(f.read())
+        self.context = self.engine.create_execution_context()
+
+    def allocate_memory(self, batch):
+        self.keypoints = np.empty(
+            (self.batch_size, self.max_keypoints, 2), dtype=np.float32
+        )
+        self.scores = np.empty((self.batch_size, self.max_keypoints), dtype=np.float32)
+        self.descriptors = np.empty(
+            (self.batch_size, self.max_keypoints, self.descriptor_dim), dtype=np.float32
+        )
+
+        # Allocate device memory
+        self.d_input = cuda.mem_alloc(1 * batch.nbytes)
+        self.d_keypoints = cuda.mem_alloc(1 * self.keypoints.nbytes)
+        self.d_scores = cuda.mem_alloc(1 * self.scores.nbytes)
+        self.d_descriptors = cuda.mem_alloc(1 * self.descriptors.nbytes)
+
+        tensor_names = [
+            self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)
+        ]
+
+        self.context.set_tensor_address(tensor_names[0], int(self.d_input))
+        self.context.set_tensor_address(tensor_names[1], int(self.d_keypoints))
+        self.context.set_tensor_address(tensor_names[2], int(self.d_scores))
+        self.context.set_tensor_address(tensor_names[3], int(self.d_descriptors))
+
+        self.stream = cuda.Stream()
+
+    def predict(self, batch):
+        if self.stream is None:
+            self.allocate_memory(batch)
+
+        # Transfer input data to device
+        cuda.memcpy_htod_async(self.d_input, batch, self.stream)
+        # Execute model
+        self.context.execute_async_v3(self.stream.handle)
+        # Transfer predictions back
+        cuda.memcpy_dtoh_async(self.keypoints, self.d_keypoints, self.stream)
+        cuda.memcpy_dtoh_async(self.scores, self.d_scores, self.stream)
+        cuda.memcpy_dtoh_async(self.descriptors, self.d_descriptors, self.stream)
+        # Syncronize threads
+        self.stream.synchronize()
+
+        return self.keypoints, self.scores, self.descriptors

src/superpoint_pruning/evaluation/plot_keypoints.py

@@ -0,0 +1,197 @@
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from superpoint_pruning.distillation.utils import load_grayscale_image, rescale_image
+from superpoint_pruning.models.superpoint import SuperPoint
+from superpoint_pruning.paths import DEFAULT_IMAGE_DIR
+from superpoint_pruning.pruning import add_pruning_parser_args, pruning_config_from_args
+
+SHARED_COLOR = "lime"
+ORIGINAL_ONLY_COLOR = "deepskyblue"
+PRUNED_ONLY_COLOR = "orangered"
+
+
+def resolve_image_path(image_dir: Path, image_name: Path) -> Path:
+    if image_name.is_absolute():
+        return image_name
+    return image_dir / image_name
+
+
+def classify_keypoints(
+    original: torch.Tensor,
+    pruned: torch.Tensor,
+    image_width: int,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    original = original.detach().cpu().round().to(torch.int64)
+    pruned = pruned.detach().cpu().round().to(torch.int64)
+    original_ids = original[:, 1] * image_width + original[:, 0]
+    pruned_ids = pruned[:, 1] * image_width + pruned[:, 0]
+
+    shared_mask = torch.isin(original_ids, pruned_ids)
+    original_only_mask = ~shared_mask
+    pruned_only_mask = ~torch.isin(pruned_ids, original_ids)
+
+    return (
+        original[shared_mask].numpy(),
+        original[original_only_mask].numpy(),
+        pruned[pruned_only_mask].numpy(),
+    )
+
+
+@torch.inference_mode()
+def extract_keypoints(
+    model: SuperPoint,
+    image: torch.Tensor,
+    scale: torch.Tensor,
+) -> torch.Tensor:
+    keypoints, _, _ = model(image)
+    keypoints = (keypoints.to(torch.float32) + 0.5) / scale - 0.5
+    return keypoints[0].detach().cpu()
+
+
+def _scatter(ax, points: np.ndarray, color: str, label: str, size: int) -> None:
+    if len(points) == 0:
+        return
+    ax.scatter(points[:, 0], points[:, 1], s=size, marker=".", color=color, label=label)
+
+
+def plot_side_by_side(
+    image: np.ndarray,
+    original: np.ndarray,
+    pruned: np.ndarray,
+    output: Path,
+    keypoint_size: int,
+) -> None:
+    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+    for ax, points, title in (
+        (axes[0], original, f"Original ({len(original)})"),
+        (axes[1], pruned, f"Pruned ({len(pruned)})"),
+    ):
+        ax.imshow(image, cmap="gray")
+        _scatter(ax, points, SHARED_COLOR, title, keypoint_size)
+        ax.set_title(title)
+        ax.axis("off")
+    fig.tight_layout()
+    output.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_overlay(
+    image: np.ndarray,
+    original: torch.Tensor,
+    pruned: torch.Tensor,
+    output: Path,
+    keypoint_size: int,
+) -> None:
+    height, width = image.shape[:2]
+    shared, original_only, pruned_only = classify_keypoints(original, pruned, width)
+    fig, ax = plt.subplots(figsize=(8, 6))
+    ax.imshow(image, cmap="gray")
+    _scatter(ax, shared, SHARED_COLOR, f"Shared ({len(shared)})", keypoint_size)
+    _scatter(
+        ax,
+        original_only,
+        ORIGINAL_ONLY_COLOR,
+        f"Original only ({len(original_only)})",
+        keypoint_size,
+    )
+    _scatter(
+        ax,
+        pruned_only,
+        PRUNED_ONLY_COLOR,
+        f"Pruned only ({len(pruned_only)})",
+        keypoint_size,
+    )
+    ax.set_title("Shared vs unique keypoints")
+    ax.axis("off")
+    ax.legend(loc="upper right", framealpha=0.8)
+    fig.tight_layout()
+    output.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+
+
+def add_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--image-dir", type=Path, default=DEFAULT_IMAGE_DIR)
+    parser.add_argument(
+        "--image-name",
+        type=Path,
+        required=True,
+        help="Image filename or absolute path.",
+    )
+    parser.add_argument("--output", type=Path, default=None, help="Output figure path.")
+    parser.add_argument(
+        "--overlay",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Plot both models on one image: shared, original-only, and pruned-only keypoints.",
+    )
+    parser.add_argument("--num-keypoints", type=int, default=512)
+    parser.add_argument("--width", type=int, default=640)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument(
+        "--hierarchical", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--skip-refinement", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument("--keypoint-size", type=int, default=5)
+    add_pruning_parser_args(parser)
+
+
+def main(args: argparse.Namespace) -> None:
+    image_path = resolve_image_path(args.image_dir, args.image_name)
+    output = args.output
+    if output is None:
+        suffix = "overlay" if args.overlay else "keypoints"
+        output = Path(f"{image_path.stem}_{suffix}.png")
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pruning_config, pruning_checkpoint = pruning_config_from_args(args)
+
+    original_image = load_grayscale_image(str(image_path))
+    image, scale = rescale_image(original_image, new_size=(args.width, args.height))
+    scale = torch.as_tensor(scale, device=device, dtype=torch.float32)
+    batch = torch.from_numpy(image[None, None].astype(np.float32)).to(device)
+
+    sp_original = SuperPoint(num_keypoints=args.num_keypoints).eval().to(device)
+    sp_pruned = SuperPoint(
+        num_keypoints=args.num_keypoints,
+        hierarchical_topk=args.hierarchical,
+        skip_refinement=args.skip_refinement,
+    )
+    if pruning_config:
+        sp_pruned.prune_backbone(pruning_config)
+    if pruning_checkpoint is not None:
+        sp_pruned.load_pruned_weights(str(pruning_checkpoint))
+    sp_pruned = sp_pruned.to(device).eval()
+
+    original_kpts = extract_keypoints(sp_original, batch, scale)
+    pruned_kpts = extract_keypoints(sp_pruned, batch, scale)
+
+    if args.overlay:
+        plot_overlay(
+            original_image, original_kpts, pruned_kpts, output, args.keypoint_size
+        )
+    else:
+        plot_side_by_side(
+            original_image,
+            original_kpts.numpy(),
+            pruned_kpts.numpy(),
+            output,
+            args.keypoint_size,
+        )
+    print(f"Saved keypoint plot to {output.resolve()}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Plot original and pruned SuperPoint keypoints."
+    )
+    add_parser_args(parser)
+    main(parser.parse_args())

src/superpoint_pruning/export.py

@@ -0,0 +1,111 @@
+"""Export SuperPoint to ONNX.
+
+Examples:
+    superpoint-pruning export --output superpoint.onnx
+    superpoint-pruning export --backbone_0_1 32 --backbone_1_0 48 --output pruned.onnx
+    superpoint-pruning export --pruning-config 32_48_64.ckpt --output pruned.onnx
+    python -m superpoint_pruning.export --output superpoint.onnx
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import torch
+from superpoint_pruning.models.superpoint import SuperPoint
+from superpoint_pruning.pruning import add_pruning_parser_args, pruning_config_from_args
+
+
+def validate_export_args(args: argparse.Namespace) -> None:
+    if args.width <= 0 or args.height <= 0 or args.batch_dim <= 0:
+        raise ValueError("Width, height, and batch_dim must be positive.")
+    if args.width % 8 or args.height % 8:
+        raise ValueError("Width and height must be divisible by 8.")
+    if args.num_keypoints <= 0 or args.num_keypoints > args.width * args.height:
+        raise ValueError("num_keypoints must be between 1 and width * height.")
+    if args.hierarchical_topk:
+        if args.hierarchical_tile_size <= 0:
+            raise ValueError("hierarchical_tile_size must be positive.")
+        if args.height % args.hierarchical_tile_size:
+            raise ValueError(
+                "hierarchical_tile_size must divide the input height when hierarchical_topk is enabled."
+            )
+        if args.num_keypoints > args.hierarchical_tile_size * args.width:
+            raise ValueError(
+                "num_keypoints must not exceed hierarchical_tile_size * width when hierarchical_topk is enabled."
+            )
+
+
+def add_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--num-keypoints", type=int, default=512)
+    parser.add_argument(
+        "--skip-refinement",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Skip the iterative NMS refinement.",
+    )
+    parser.add_argument(
+        "--hierarchical-topk",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument("--hierarchical-tile-size", type=int, default=32)
+    parser.add_argument("--width", type=int, default=640)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--batch-dim", type=int, default=1)
+    parser.add_argument("--output", type=Path, default=Path("superpoint.onnx"))
+    add_pruning_parser_args(parser)
+    parser.add_argument(
+        "--skip-bn", action=argparse.BooleanOptionalAction, default=False
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Export a SuperPoint model to ONNX.")
+    add_parser_args(parser)
+    return parser
+
+
+def main(args: argparse.Namespace) -> None:
+    pruning_config, pruning_checkpoint = pruning_config_from_args(args)
+    validate_export_args(args)
+
+    print("num_keypoints: ", args.num_keypoints)
+    print("skip_refinement: ", args.skip_refinement)
+    print("hierarchical_topk: ", args.hierarchical_topk)
+    print("hierarchical_tile_size: ", args.hierarchical_tile_size)
+    print("width: ", args.width)
+    print("height: ", args.height)
+    print("batch_dim: ", args.batch_dim)
+    print("output: ", args.output)
+    print("pruning_config: ", args.pruning_config)
+    print("skip_bn: ", args.skip_bn)
+
+    model = SuperPoint(
+        num_keypoints=args.num_keypoints,
+        skip_refinement=args.skip_refinement,
+        hierarchical_topk=args.hierarchical_topk,
+        hierarchical_tile_size=args.hierarchical_tile_size,
+        use_bn=not args.skip_bn,
+    )
+    if pruning_config:
+        model.prune_backbone(pruning_config)
+    if pruning_checkpoint:
+        model.load_pruned_weights(str(pruning_checkpoint))
+
+    inputs = torch.zeros(args.batch_dim, 1, args.height, args.width)
+    model.eval()
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with torch.no_grad():
+        torch.onnx.export(
+            model.cpu(),
+            inputs,
+            args.output,
+            input_names=["inputs"],
+            output_names=["keypoints", "scores", "descriptors"],
+        )
+    print(f"Exported {args.output} with input shape {tuple(inputs.shape)}.")
+
+
+if __name__ == "__main__":
+    main(build_parser().parse_args())

src/superpoint_pruning/models/superpoint.py

@@ -0,0 +1,354 @@
+"""PyTorch implementation of the SuperPoint model,
+   derived from the TensorFlow re-implementation (2018).
+   Authors: Rémi Pautrat, Paul-Edouard Sarlin
+"""
+
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+
+from superpoint_pruning.paths import DEFAULT_WEIGHTS_PATH
+
+V6_PREFIXES = [
+    "backbone.0.0",
+    "backbone.0.1",
+    "backbone.1.0",
+    "backbone.1.1",
+    "backbone.2.0",
+    "backbone.2.1",
+    "backbone.3.0",
+    "backbone.3.1",
+    "detector.0",
+    "detector.1",
+    "descriptor.0",
+    "descriptor.1",
+]
+
+
+def convert_v6_state_dict(state_dict: dict) -> dict:
+    converted = {}
+    for key, value in state_dict.items():
+        for prefix in V6_PREFIXES:
+            token = prefix + "."
+            if not key.startswith(token):
+                continue
+            kind, param = key[len(token) :].split(".", 1)
+            if kind == "conv":
+                converted[f"{prefix.replace('.', '_')}.{param}"] = value
+            elif kind == "bn":
+                converted[f"{prefix.replace('.', '_')}_bn.{param}"] = value
+            else:
+                raise KeyError(f"Unrecognized v6 submodule in '{key}'")
+            break
+        else:
+            raise KeyError(f"Unrecognized v6 key: {key}")
+    return converted
+
+
+def sample_descriptors(keypoints, descriptors, s: int = 8):
+    b, c, h, w = descriptors.shape
+    divisor = (
+        torch._shape_as_tensor(descriptors)[[3, 2]]
+        .to(keypoints.dtype)
+        .to(keypoints.device)
+        * s
+    )
+    keypoints = (keypoints + 0.5) / divisor
+    keypoints = keypoints * 2 - 1
+    descriptors = torch.nn.functional.grid_sample(
+        descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", align_corners=False
+    )
+    descriptors = torch.nn.functional.normalize(
+        descriptors.reshape(b, c, -1), p=2, dim=1
+    ).permute(0, 2, 1)
+    return descriptors
+
+
+def batched_nms(scores, nms_radius: int, skip_refinement: bool = False):
+    assert nms_radius >= 0
+
+    def max_pool(x):
+        return torch.nn.functional.max_pool2d(
+            x, kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius
+        )
+
+    scores = scores[:, None]
+    zeros = torch.zeros_like(scores)
+    max_mask = scores == max_pool(scores)
+    if not skip_refinement:
+        for _ in range(2):
+            supp_mask = max_pool(max_mask.float()) > 0
+            supp_scores = torch.where(supp_mask, zeros, scores)
+            new_max_mask = supp_scores == max_pool(supp_scores)
+            max_mask = max_mask | (new_max_mask & (~supp_mask))
+    return torch.where(max_mask, scores, zeros)[:, 0]
+
+
+def hierarchical_topk(scores, tile_size: int, num_keypoints: int):
+    B, H, W = scores.shape
+    tile_h = tile_size
+    assert H % tile_h == 0, "Tile size must divide the height of the scores"
+    scores_tiled = scores.reshape(B, H // tile_h, tile_h * W)
+
+    local_scores, local_idx = scores_tiled.topk(num_keypoints, dim=-1, sorted=False)
+
+    # Convert local indices into global flattened indices.
+    tile_offset = (
+        torch.arange(H // tile_h, device=scores.device).view(1, -1, 1) * tile_h * W
+    )
+    global_idx = local_idx + tile_offset
+
+    local_scores = local_scores.reshape(B, -1)
+    global_idx = global_idx.reshape(B, -1)
+
+    top_scores, sel = local_scores.topk(num_keypoints, dim=-1, sorted=True)
+    top_indices = global_idx.gather(1, sel)
+
+    return top_scores, top_indices
+
+
+def get_conv_layer(c_in, c_out, kernel_size, relu=True):
+    padding = (kernel_size - 1) // 2
+    conv = nn.Conv2d(c_in, c_out, kernel_size=kernel_size, stride=1, padding=padding)
+    return conv
+
+
+class SuperPoint(nn.Module):
+    default_conf = {
+        "nms_radius": 4,
+        "num_keypoints": 1024,
+        "remove_borders": 4,
+        "descriptor_dim": 256,
+        "channels": [64, 64, 128, 128, 256],
+        "skip_refinement": False,
+        "hierarchical_topk": False,
+        "hierarchical_tile_size": 32,
+        "use_bn": True,
+        "default_weights_path": DEFAULT_WEIGHTS_PATH,
+        "load_default_weights": True,
+        "return_dense": False,
+    }
+
+    def __init__(self, **conf):
+        super().__init__()
+        conf = {**self.default_conf, **conf}
+        self.conf = SimpleNamespace(**conf)
+        self.stride = 2 ** (len(self.conf.channels) - 2)
+        channels = [1, *self.conf.channels[:-1]]
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.relu = nn.ReLU(inplace=True)
+
+        for i, c in enumerate(channels[1:]):
+            self.add_module(f"backbone_{i}_0", get_conv_layer(channels[i], c, 3))
+            if self.conf.use_bn:
+                self.add_module(f"backbone_{i}_0_bn", nn.BatchNorm2d(c, eps=0.001))
+            self.add_module(f"backbone_{i}_1", get_conv_layer(c, c, 3))
+            if self.conf.use_bn:
+                self.add_module(f"backbone_{i}_1_bn", nn.BatchNorm2d(c, eps=0.001))
+
+        c = self.conf.channels[-1]
+        self.add_module("detector_0", get_conv_layer(channels[-1], c, 3))
+        if self.conf.use_bn:
+            self.add_module("detector_0_bn", nn.BatchNorm2d(c, eps=0.001))
+        self.add_module("detector_1", get_conv_layer(c, self.stride**2 + 1, 1))
+        if self.conf.use_bn:
+            self.add_module(
+                "detector_1_bn", nn.BatchNorm2d(self.stride**2 + 1, eps=0.001)
+            )
+        self.add_module("descriptor_0", get_conv_layer(channels[-1], c, 3))
+        if self.conf.use_bn:
+            self.add_module("descriptor_0_bn", nn.BatchNorm2d(c, eps=0.001))
+        self.add_module("descriptor_1", get_conv_layer(c, self.conf.descriptor_dim, 1))
+        if self.conf.use_bn:
+            self.add_module(
+                "descriptor_1_bn", nn.BatchNorm2d(self.conf.descriptor_dim, eps=0.001)
+            )
+
+        if self.conf.use_bn and self.conf.load_default_weights:
+            self.load_default_weights(self.conf.default_weights_path)
+
+    def _forward_conv(
+        self, name: str, x: torch.Tensor, relu: bool = True
+    ) -> torch.Tensor:
+        x = getattr(self, name)(x)
+        if relu:
+            x = self.relu(x)
+        if self.conf.use_bn:
+            x = getattr(self, name + "_bn")(x)
+        return x
+
+    def load_default_weights(self, weights_path: str) -> None:
+        """Load and rename superpoint_v6_from_tf.pth weights."""
+        checkpoint = torch.load(weights_path, map_location="cpu")
+        state_dict = convert_v6_state_dict(checkpoint)
+        missing, unexpected = self.load_state_dict(state_dict, strict=False)
+        missing = [k for k in missing if not k.endswith("num_batches_tracked")]
+        if missing or unexpected:
+            raise RuntimeError(
+                f"Failed to load weights from {weights_path}. missing={missing} unexpected={unexpected}"
+            )
+        print(f"Loaded default weights from {weights_path}")
+
+    def load_pruned_weights(self, checkpoint_path: str) -> None:
+        """Load pruned SP weights from checkpoint"""
+        checkpoint = torch.load(checkpoint_path, map_location="cpu")["state_dict"]
+        checkpoint = {k.replace("model.", ""): v for k, v in checkpoint.items()}
+        self.load_state_dict(checkpoint, strict=True)
+
+    def _bn_name(self, conv_name: str) -> str:
+        return conv_name + "_bn"
+
+    def _previous_backbone_layer(self, layer: str) -> str:
+        parts = layer.split("_")
+        stage, index = int(parts[1]), int(parts[2])
+        if index == 1:
+            return f"backbone_{stage}_0"
+        if index == 0 and stage > 0:
+            return f"backbone_{stage - 1}_1"
+        raise ValueError(f"Invalid layer: {layer}")
+
+    def _prune_bn(self, conv_name: str, keep_idx: torch.Tensor) -> None:
+        """Resize the BN that follows a pruned conv so channel counts still match."""
+        bn_name = self._bn_name(conv_name)
+        old_bn = getattr(self, bn_name)
+        new_bn = nn.BatchNorm2d(
+            len(keep_idx),
+            eps=old_bn.eps,
+            momentum=old_bn.momentum,
+            affine=old_bn.affine,
+            track_running_stats=old_bn.track_running_stats,
+        )
+        with torch.no_grad():
+            if old_bn.affine:
+                new_bn.weight.copy_(old_bn.weight[keep_idx])
+                new_bn.bias.copy_(old_bn.bias[keep_idx])
+            if old_bn.track_running_stats:
+                new_bn.running_mean.copy_(old_bn.running_mean[keep_idx])
+                new_bn.running_var.copy_(old_bn.running_var[keep_idx])
+                new_bn.num_batches_tracked.copy_(old_bn.num_batches_tracked)
+        setattr(self, bn_name, new_bn)
+
+    def prune_backbone(self, config: dict):
+        for layer, channel in config.items():
+            previous_layer = self._previous_backbone_layer(layer)
+            old1 = getattr(self, previous_layer)
+            old2 = getattr(self, layer)
+            magnitude = old2.weight.abs().mean(dim=(0, 2, 3))
+            keep_idx = (
+                torch.topk(magnitude, k=channel, largest=True).indices.sort().values
+            )
+            new1 = torch.nn.Conv2d(
+                old1.in_channels,
+                channel,
+                kernel_size=old1.kernel_size,
+                stride=1,
+                padding=1,
+            )
+            new2 = torch.nn.Conv2d(
+                channel,
+                old2.out_channels,
+                kernel_size=old2.kernel_size,
+                stride=1,
+                padding=1,
+            )
+            with torch.no_grad():
+                new1.weight.copy_(old1.weight[keep_idx, :, :, :])
+                new1.bias.copy_(old1.bias[keep_idx])
+                new2.weight.copy_(old2.weight[:, keep_idx, :, :])
+                new2.bias.copy_(old2.bias)
+            setattr(self, previous_layer, new1)
+            setattr(self, layer, new2)
+            if self.conf.use_bn:
+                self._prune_bn(previous_layer, keep_idx)
+
+        print("Pruned model structure:")
+        print(self)
+
+    def dense_head(
+        self, scores: torch.Tensor, descriptors_dense: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+        descriptors_dense = torch.nn.functional.normalize(descriptors_dense, p=2, dim=1)
+
+        scores = torch.nn.functional.softmax(scores, 1)[:, :-1]
+        b, _, h, w = scores.shape
+        scores = scores.permute(0, 2, 3, 1).reshape(b, h, w, self.stride, self.stride)
+        scores = scores.permute(0, 1, 3, 2, 4).reshape(
+            b, h * self.stride, w * self.stride
+        )
+        scores = batched_nms(
+            scores, self.conf.nms_radius, skip_refinement=self.conf.skip_refinement
+        )
+
+        # Discard keypoints near the image borders
+        if self.conf.remove_borders:
+            pad = self.conf.remove_borders
+            scores[:, :pad] = -1
+            scores[:, :, :pad] = -1
+            scores[:, -pad:] = -1
+            scores[:, :, -pad:] = -1
+
+        if self.conf.hierarchical_topk:
+            top_scores, top_indices = hierarchical_topk(
+                scores, self.conf.hierarchical_tile_size, self.conf.num_keypoints
+            )
+        else:
+            top_scores, top_indices = scores.reshape(
+                b, h * self.stride * w * self.stride
+            ).topk(self.conf.num_keypoints)
+        y_idx = torch.div(top_indices, w * self.stride, rounding_mode="floor")
+        x_idx = torch.remainder(top_indices, w * self.stride)
+        top_keypoints = torch.stack((x_idx, y_idx), dim=-1).to(dtype=torch.float32)
+        top_descriptors = sample_descriptors(
+            top_keypoints, descriptors_dense, self.stride
+        )
+
+        return (
+            top_keypoints,
+            top_scores,
+            top_descriptors,
+        )
+
+    def forward(
+        self, image: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+        x = self._forward_conv("backbone_0_0", image)
+        x = self._forward_conv("backbone_0_1", x)
+        x = self.pool(x)
+        x = self._forward_conv("backbone_1_0", x)
+        x = self._forward_conv("backbone_1_1", x)
+        x = self.pool(x)
+        x = self._forward_conv("backbone_2_0", x)
+        x = self._forward_conv("backbone_2_1", x)
+        x = self.pool(x)
+        x = self._forward_conv("backbone_3_0", x)
+        x = self._forward_conv("backbone_3_1", x)
+        descriptors_dense = self._forward_conv(
+            "descriptor_1", self._forward_conv("descriptor_0", x), relu=False
+        )
+        scores = self._forward_conv(
+            "detector_1", self._forward_conv("detector_0", x), relu=False
+        )
+
+        if self.conf.return_dense:
+            return scores, descriptors_dense
+
+        return self.dense_head(scores, descriptors_dense)
+
+
+if __name__ == "__main__":
+    sp = SuperPoint(num_keypoints=512)
+    inputs = torch.zeros(1, 1, 768, 1024)
+    torch.onnx.export(
+        sp.cpu(),
+        inputs,
+        f"SP.onnx",
+        input_names=["inputs"],
+        output_names=["keypoints", "scores", "descriptors"],
+        # opset_version=opset,
+        # dynamic_axes=dynamic_axes,
+        # dynamic_shapes=dynamic_shapes,
+        # dynamo=True,
+    )

src/superpoint_pruning/paths.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+PACKAGE_ROOT = Path(__file__).resolve().parent
+DATA_ROOT = PACKAGE_ROOT / "distillation" / "data"
+DEFAULT_CONFIG_PATH = PACKAGE_ROOT / "distillation" / "base_config.yaml"
+DEFAULT_WEIGHTS_PATH = PACKAGE_ROOT / "weights" / "superpoint_v6_from_tf.pth"
+DEFAULT_DATASET_NAME = "indoor_forward_3_snapdragon_with_gt"
+DEFAULT_IMAGE_DIR = DATA_ROOT / "datasets" / DEFAULT_DATASET_NAME / "img"
+DEFAULT_TRAINING_IDS = DATA_ROOT / "gt_id_file_train_base_250.txt"

src/superpoint_pruning/pruning.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from superpoint_pruning.models.superpoint import SuperPoint
+
+PRUNABLE_LAYERS = (
+    "backbone_0_1",
+    "backbone_1_0",
+    "backbone_1_1",
+    "backbone_2_0",
+    "backbone_2_1",
+    "backbone_3_0",
+    "backbone_3_1",
+)
+_C = SuperPoint.default_conf["channels"]
+MAX_PRUNED_CHANNELS = {
+    "backbone_0_1": _C[0],
+    "backbone_1_0": _C[0],
+    "backbone_1_1": _C[1],
+    "backbone_2_0": _C[1],
+    "backbone_2_1": _C[2],
+    "backbone_3_0": _C[2],
+    "backbone_3_1": _C[3],
+}
+
+
+def parse_pruning_config(checkpoint_path: str) -> dict[str, int]:
+    """Build a pruning config from a checkpoint named ``a_b_c...ckpt``."""
+    checkpoint = Path(checkpoint_path)
+    if checkpoint.suffix != ".ckpt":
+        raise ValueError("A pruning checkpoint must have a .ckpt extension.")
+
+    values = checkpoint.stem.split("_")
+    if not 1 <= len(values) <= len(PRUNABLE_LAYERS):
+        raise ValueError(
+            f"Expected 1-{len(PRUNABLE_LAYERS)} underscore-separated channel counts in '{checkpoint.name}'."
+        )
+
+    try:
+        config = {layer: int(value) for layer, value in zip(PRUNABLE_LAYERS, values)}
+    except ValueError as error:
+        raise ValueError(
+            "Pruning checkpoint names must contain only integer channel counts, for example '32_48_64.ckpt'."
+        ) from error
+
+    if any(channels <= 0 for channels in config.values()):
+        raise ValueError("Pruned channel counts must be positive.")
+    return config
+
+
+def validate_pruning_config(pruning_config: dict[str, int]) -> None:
+    for layer, channels in pruning_config.items():
+        max_channels = MAX_PRUNED_CHANNELS[layer]
+        if channels > max_channels:
+            raise ValueError(
+                f"{layer} cannot be pruned to {channels} channels; it has only {max_channels} input channels."
+            )
+
+
+def add_pruning_parser_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--pruning-config",
+        type=Path,
+        help=(
+            "Checkpoint path with pruned weights. Its basename must be named like '32_48_64.ckpt'; "
+            "counts map to backbone_0_1, backbone_1_0, backbone_1_1, and so on."
+        ),
+    )
+    for layer in PRUNABLE_LAYERS:
+        parser.add_argument(
+            f"--{layer}",
+            dest=layer,
+            type=int,
+            metavar="CHANNELS",
+            help=f"Prune the input channels of {layer}.",
+        )
+
+
+def pruning_config_from_args(
+    args: argparse.Namespace,
+) -> tuple[dict[str, int], Path | None]:
+    cli_pruning_config = {
+        layer: getattr(args, layer)
+        for layer in PRUNABLE_LAYERS
+        if getattr(args, layer) is not None
+    }
+    if args.pruning_config and cli_pruning_config:
+        raise ValueError(
+            "Use either --pruning-config or individual layer options, not both."
+        )
+
+    pruning_config = (
+        parse_pruning_config(str(args.pruning_config))
+        if args.pruning_config
+        else cli_pruning_config
+    )
+    validate_pruning_config(pruning_config)
+    return pruning_config, args.pruning_config

src/superpoint_pruning/weights/16_16_24_32_64.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb893fb8121e2881d41132510ed6ff1708004f1c78a73f3dcf4fa2b237e7ea2d
+size 12994600

src/superpoint_pruning/weights/superpoint_v6_from_tf.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd5d19a5061848e248c17728878ea166b66512076d43c77dbcf27f4a88a56084
+size 5251225