Spaces:
Sleeping
Sleeping
| 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 | |