|
import cv2 |
|
import numpy as np |
|
import torch |
|
import os |
|
|
|
import sys |
|
|
|
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) |
|
sys.path.insert(0, ROOT_DIR) |
|
|
|
from superpoint import SuperPoint |
|
|
|
|
|
def resize(img, resize): |
|
img_h, img_w = img.shape[0], img.shape[1] |
|
cur_size = max(img_h, img_w) |
|
if len(resize) == 1: |
|
scale1, scale2 = resize[0] / cur_size, resize[0] / cur_size |
|
else: |
|
scale1, scale2 = resize[0] / img_h, resize[1] / img_w |
|
new_h, new_w = int(img_h * scale1), int(img_w * scale2) |
|
new_img = cv2.resize(img.astype("float32"), (new_w, new_h)).astype("uint8") |
|
scale = np.asarray([scale2, scale1]) |
|
return new_img, scale |
|
|
|
|
|
class ExtractSIFT: |
|
def __init__(self, config, root=True): |
|
self.num_kp = config["num_kpt"] |
|
self.contrastThreshold = config["det_th"] |
|
self.resize = config["resize"] |
|
self.root = root |
|
|
|
def run(self, img_path): |
|
self.sift = cv2.xfeatures2d.SIFT_create( |
|
nfeatures=self.num_kp, contrastThreshold=self.contrastThreshold |
|
) |
|
img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) |
|
scale = [1, 1] |
|
if self.resize[0] != -1: |
|
img, scale = resize(img, self.resize) |
|
cv_kp, desc = self.sift.detectAndCompute(img, None) |
|
kp = np.array( |
|
[ |
|
[_kp.pt[0] / scale[1], _kp.pt[1] / scale[0], _kp.response] |
|
for _kp in cv_kp |
|
] |
|
) |
|
index = np.flip(np.argsort(kp[:, 2])) |
|
kp, desc = kp[index], desc[index] |
|
if self.root: |
|
desc = np.sqrt( |
|
abs(desc / (np.linalg.norm(desc, axis=-1, ord=1)[:, np.newaxis] + 1e-8)) |
|
) |
|
return kp[: self.num_kp], desc[: self.num_kp] |
|
|
|
|
|
class ExtractSuperpoint(object): |
|
def __init__(self, config): |
|
default_config = { |
|
"descriptor_dim": 256, |
|
"nms_radius": 4, |
|
"detection_threshold": config["det_th"], |
|
"max_keypoints": config["num_kpt"], |
|
"remove_borders": 4, |
|
"model_path": "../weights/sp/superpoint_v1.pth", |
|
} |
|
self.superpoint_extractor = SuperPoint(default_config) |
|
self.superpoint_extractor.eval(), self.superpoint_extractor.cuda() |
|
self.num_kp = config["num_kpt"] |
|
if "padding" in config.keys(): |
|
self.padding = config["padding"] |
|
else: |
|
self.padding = False |
|
self.resize = config["resize"] |
|
|
|
def run(self, img_path): |
|
img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) |
|
scale = 1 |
|
if self.resize[0] != -1: |
|
img, scale = resize(img, self.resize) |
|
with torch.no_grad(): |
|
result = self.superpoint_extractor( |
|
torch.from_numpy(img / 255.0).float()[None, None].cuda() |
|
) |
|
score, kpt, desc = ( |
|
result["scores"][0], |
|
result["keypoints"][0], |
|
result["descriptors"][0], |
|
) |
|
score, kpt, desc = score.cpu().numpy(), kpt.cpu().numpy(), desc.cpu().numpy().T |
|
kpt = np.concatenate([kpt / scale, score[:, np.newaxis]], axis=-1) |
|
|
|
if self.padding: |
|
if len(kpt) < self.num_kp: |
|
res = int(self.num_kp - len(kpt)) |
|
pad_x, pad_desc = np.random.uniform(size=[res, 2]) * ( |
|
img.shape[0] + img.shape[1] |
|
) / 2, np.random.uniform(size=[res, 256]) |
|
pad_kpt, pad_desc = ( |
|
np.concatenate([pad_x, np.zeros([res, 1])], axis=-1), |
|
pad_desc / np.linalg.norm(pad_desc, axis=-1)[:, np.newaxis], |
|
) |
|
kpt, desc = np.concatenate([kpt, pad_kpt], axis=0), np.concatenate( |
|
[desc, pad_desc], axis=0 |
|
) |
|
return kpt, desc |
|
|
