Magic Leap pretrained model in classical baselines + adaptation of cl…

…assical methods to normalized images
lzb863 · Nov 3, 2019 · bca50ff · bca50ff
1 parent 019c467
commit bca50ff
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/superpoint/models/classical_detectors.py b/superpoint/models/classical_detectors.py
@@ -1,12 +1,17 @@
 import tensorflow as tf
 import numpy as np
 import cv2
+import sys
+
+sys.path.append('/cluster/home/pautratr/3d_project/SuperPointPretrainedNetwork')
 
 from .base_model import BaseModel
 from .utils import box_nms
+from demo_superpoint import SuperPointNet, SuperPointFrontend
 
 
 def classical_detector(im, **config):
+    im = np.uint8(im * 255)
     if config['method'] == 'harris':
         detections = cv2.cornerHarris(im, 4, 3, 0.04)
 
@@ -29,6 +34,15 @@ def classical_detector(im, **config):
     elif config['method'] == 'random':
         detections = np.random.rand(im.shape[0], im.shape[1])
 
+    elif config['method'] == 'pretrained_magic_point':
+        weights_path = '/cluster/home/pautratr/3d_project/SuperPointPretrainedNetwork/superpoint_v1.pth'
+        fe = SuperPointFrontend(weights_path=weights_path,
+                                nms_dist=config['nms'],
+                                conf_thresh=0.015,
+                                nn_thresh=0.7,
+                                cuda=True)
+        points, desc, detections = fe.run(im[:, :, 0])
+
     return detections.astype(np.float32)
 
 

diff --git a/superpoint/models/classical_detectors_descriptors.py b/superpoint/models/classical_detectors_descriptors.py
@@ -1,14 +1,18 @@
 import tensorflow as tf
 import numpy as np
 import cv2
+import sys
+
+sys.path.append('/cluster/home/pautratr/3d_project/SuperPointPretrainedNetwork')
 
 from .base_model import BaseModel
 from .utils import box_nms
+from demo_superpoint import SuperPointNet, SuperPointFrontend
 
 
 def classical_detector_descriptor(im, **config):
-    im = np.uint8(im)
     if config['method'] == 'sift':
+        im = np.uint8(im * 255)
         sift = cv2.xfeatures2d.SIFT_create(nfeatures=1500)
         keypoints, desc = sift.detectAndCompute(im, None)
         responses = np.array([k.response for k in keypoints])
@@ -21,6 +25,7 @@ def classical_detector_descriptor(im, **config):
         descriptors[keypoints[:, 1], keypoints[:, 0]] = desc
 
     elif config['method'] == 'orb':
+        im = np.uint8(im * 255)
         orb = cv2.ORB_create(nfeatures=1500)
         keypoints, desc = orb.detectAndCompute(im, None)
         responses = np.array([k.response for k in keypoints])
@@ -32,6 +37,18 @@ def classical_detector_descriptor(im, **config):
         descriptors = np.zeros((im.shape[0], im.shape[1], 32), np.float)
         descriptors[keypoints[:, 1], keypoints[:, 0]] = desc
 
+    elif config['method'] == 'pretrained_magic_point':
+        weights_path = '/cluster/home/pautratr/3d_project/SuperPointPretrainedNetwork/superpoint_v1.pth'
+        fe = SuperPointFrontend(weights_path=weights_path,
+                                nms_dist=config['nms'],
+                                conf_thresh=0.015,
+                                nn_thresh=0.7,
+                                cuda=False)
+        points, desc, detections = fe.run(im[:, :, 0])
+        points = points.astype(int)
+        descriptors = np.zeros((im.shape[0], im.shape[1], 256), np.float)
+        descriptors[points[1, :], points[0, :]] = np.transpose(desc)
+
     detections = detections.astype(np.float32)
     descriptors = descriptors.astype(np.float32)
     return (detections, descriptors)