|
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.).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 |
|
|
