def __init__(self, do_cuda=True): print('Using TfeatFeature2D') self.model_base_path = config.cfg.root_folder + '/thirdparty/tfeat/' self.do_cuda = do_cuda & torch.cuda.is_available() print('cuda:', self.do_cuda) device = torch.device("cuda:0" if self.do_cuda else "cpu") torch.set_grad_enabled(False) # mag_factor is how many times the original keypoint scale # is enlarged to generate a patch from a keypoint self.mag_factor = 3 print('==> Loading pre-trained network.') #init tfeat and load the trained weights self.model = tfeat_model.TNet() self.models_path = self.model_base_path + 'pretrained-models' self.net_name = 'tfeat-liberty' self.model.load_state_dict( torch.load( os.path.join(self.models_path, self.net_name + ".params"))) if self.do_cuda: self.model.cuda() print('Extracting on GPU') else: print('Extracting on CPU') self.model = model.cpu() self.model.eval() print('==> Successfully loaded pre-trained network.')
import os import cv2 import tfeat_utils import numpy as np import cv2 from matplotlib import pyplot as plt tfeat_base_path = '../../thirdparty/tfeat/' # open images img1 = cv2.imread(tfeat_base_path + 'imgs/v_churchill/1.ppm', 0) img2 = cv2.imread(tfeat_base_path + 'imgs/v_churchill/6.ppm', 0) #init tfeat and load the trained weights tfeat = tfeat_model.TNet() models_path = tfeat_base_path + 'pretrained-models' net_name = 'tfeat-liberty' tfeat.load_state_dict( torch.load(os.path.join(models_path, net_name + ".params"))) tfeat.cuda() tfeat.eval() # use BRISK detector brisk = cv2.BRISK_create() kp1, des1 = brisk.detectAndCompute(img1, None) kp2, des2 = brisk.detectAndCompute(img2, None) bf = cv2.BFMatcher(cv2.NORM_HAMMING) matches = bf.knnMatch(des1, des2, k=2) # Apply ratio test
def main(): args = parse_args() tfeat = tfeat_model.TNet() tfeat.load_state_dict(torch.load(args.model_path)) tfeat.cuda() tfeat.eval() if not os.path.exists(os.path.join(args.dataset_path, "descriptors")): os.makedirs(os.path.join(args.dataset_path, "descriptors")) image_names = os.listdir(os.path.join(args.dataset_path, "images")) for i, image_name in enumerate(image_names): patches_path = os.path.join(args.dataset_path, "descriptors", image_name + ".bin.patches.mat") if not os.path.exists(patches_path): continue print("Computing features for {} [{}/{}]".format( image_name, i + 1, len(image_names)), end="") start_time = time.time() descriptors_path = os.path.join(args.dataset_path, "descriptors", image_name + ".bin") if os.path.exists(descriptors_path): print(" -> skipping, already exist") continue with h5py.File(patches_path, 'r') as patches_file: patches31 = np.array(patches_file["patches"]).T if patches31.ndim != 3: print(" -> skipping, invalid input") write_matrix(descriptors_path, np.zeros((0, 128), dtype=np.float32)) continue patches = np.empty((patches31.shape[0], 32, 32), dtype=np.float32) patches[:, :31, :31] = patches31 patches[:, 31, :31] = patches31[:, 30, :] patches[:, :31, 31] = patches31[:, :, 30] patches[:, 31, 31] = patches31[:, 30, 30] descriptors = [] for i in range(0, patches.shape[0], args.batch_size): patches_batch = \ patches[i:min(i + args.batch_size, patches.shape[0])] patches_batch = \ torch.from_numpy(patches_batch[:, None]).float().cuda() descriptors.append(tfeat(patches_batch).detach().cpu().numpy()) if len(descriptors) == 0: descriptors = np.zeros((0, 128), dtype=np.float32) else: descriptors = np.concatenate(descriptors) write_matrix(descriptors_path, descriptors) print(" in {:.3f}s".format(time.time() - start_time))