Spaces:
Running
Running
File size: 3,078 Bytes
a80d6bb |
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 |
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.).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
|