"""
    "XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
    https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/

Minimal example of how to use XFeat.
"""

import numpy as np
import os
import torch
import tqdm

from modules.xfeat import XFeat

os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Force CPU, comment for GPU

xfeat = XFeat()

# Random input
x = torch.randn(1, 3, 480, 640)

# Simple inference with batch = 1
output = xfeat.detectAndCompute(x, top_k=4096)[0]
print("----------------")
print("keypoints: ", output["keypoints"].shape)
print("descriptors: ", output["descriptors"].shape)
print("scores: ", output["scores"].shape)
print("----------------\n")

x = torch.randn(1, 3, 480, 640)
# Stress test
for i in tqdm.tqdm(range(100), desc="Stress test on VGA resolution"):
    output = xfeat.detectAndCompute(x, top_k=4096)

# Batched mode
x = torch.randn(4, 3, 480, 640)
outputs = xfeat.detectAndCompute(x, top_k=4096)
print("# detected features on each batch item:", [len(o["keypoints"]) for o in outputs])

# Match two images with sparse features
x1 = torch.randn(1, 3, 480, 640)
x2 = torch.randn(1, 3, 480, 640)
mkpts_0, mkpts_1 = xfeat.match_xfeat(x1, x2)

# Match two images with semi-dense approach -- batched mode with batch size 4
x1 = torch.randn(4, 3, 480, 640)
x2 = torch.randn(4, 3, 480, 640)
matches_list = xfeat.match_xfeat_star(x1, x2)
print(matches_list[0].shape)
