Spaces:
Running
Running
File size: 3,833 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
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
]
) # N*3
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)
# padding randomly
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
|