vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/LISRD/lisrd/export_features.py

@@ -0,0 +1,148 @@
+"""
+Export features detections and descriptions for all images in a given folder.
+"""
+
+import os
+import sys
+import argparse
+import numpy as np
+import torch
+import torch.nn.functional as func
+import cv2
+from tqdm import tqdm
+
+from .models import get_model
+from .models.base_model import Mode
+from .datasets.utils.data_reader import resize_and_crop
+from .utils.pytorch_utils import keypoints_to_grid
+from .models.keypoint_detectors import SIFT_detect, SP_detect, load_SP_net
+
+
+base_config = {
+    'learning_rate': 0.001, 'desc_size': 128,
+    'tile': 3, 'n_clusters': 8, 'meta_desc_dim': 128,
+    'compute_meta_desc': True, 'freeze_local_desc': False}
+
+
+def export(images_list, model, checkpoint, keypoints_type,
+           num_keypoints, detection_thresh, extension,
+           resize=False, h=480, w=640):
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    config = base_config
+    config['name'] = model
+    config['n_kp'] = num_keypoints
+
+    # Load the model
+    net = get_model(config['name'])(None, config, device)
+    net.load(checkpoint, Mode.EXPORT)
+    net._net.eval()
+
+    # Load the keypoint network if necessary
+    if keypoints_type == 'superpoint':
+        kp_net = load_SP_net(conf_thresh=detection_thresh)
+
+    # Parse the data, predict the features, and export them in an npz file
+    with open(images_list, 'r') as f:
+        image_files = f.readlines()
+    image_files = [path.strip('\n') for path in image_files]
+
+    for img_path in tqdm(image_files):
+        img = cv2.imread(img_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        if resize:
+            img = resize_and_crop(img, (h, w))
+        img_size = img.shape
+        if img_size[2] != 3:
+            sys.exit('Export only available for RGB images.')
+        cpu_gray_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+        img = torch.tensor(img, dtype=torch.float, device=device)
+        img = img.permute(2, 0, 1).unsqueeze(0) / 255.
+
+        # Predict keypoints
+        if keypoints_type == 'sift':
+            cpu_gray_img = np.uint8(cpu_gray_img)
+            keypoints = SIFT_detect(cpu_gray_img, nfeatures=num_keypoints,
+                                    contrastThreshold=0.04)
+        if keypoints_type == 'superpoint':
+            keypoints = SP_detect(cpu_gray_img, kp_net)
+        scores = keypoints[:, 2]
+
+        grid_points = keypoints_to_grid(
+            torch.tensor(keypoints[:, :2], dtype=torch.float, device=device),
+            img_size[:2])
+        keypoints = keypoints[:, [1, 0]]
+
+        # Predict the corresponding descriptors
+        inputs = {'image0': img}
+        with torch.no_grad():
+            outputs = net._forward(inputs, Mode.EXPORT, config)
+        
+        descs = outputs['descriptors']
+        meta_descs = outputs['meta_descriptors']
+        if model == 'lisrd':
+            descriptors, meta_descriptors = [], []
+            for k in descs.keys():
+                desc = func.normalize(
+                    func.grid_sample(descs[k], grid_points),
+                    dim=1).squeeze().cpu().numpy().transpose(1, 0)
+                descriptors.append(desc)
+                meta_descriptors.append(
+                    meta_descs[k].squeeze().cpu().numpy())
+        else:
+            descriptors = [v[0] for v in descs.values()]
+            meta_descriptors = [v[0] for v in meta_descs.values()]
+        descriptors = np.stack(descriptors, axis=1)
+        meta_descriptors = np.stack(meta_descriptors, axis=0)
+
+        # Keep the best scores
+        idxs = scores.argsort()[-num_keypoints:]
+
+        with open(img_path + extension, 'wb') as output_file:
+            np.savez(output_file, keypoints=keypoints[idxs],
+                     descriptors=descriptors[idxs], scores=scores[idxs],
+                     meta_descriptors=meta_descriptors)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('images_list', type=str,
+                         help='Path to a txt file containing the image paths.')
+    parser.add_argument('model', type=str,
+                         help="Model name.")
+    parser.add_argument('--checkpoint', type=str, default=None,
+                         help="Path to the model checkpoint.")
+    parser.add_argument('--keypoints', type=str, default='sift',
+                         help="Type of keypoints to use: "
+                         + "'sift', or 'superpoint'.")
+    parser.add_argument('--num_kp', type=int, default=2000,
+                         help="Number of keypoints to use.")
+    parser.add_argument('--detection_thresh', type=float, default=0.015,
+                         help="Detection threshold for SuperPoint.")
+    parser.add_argument('--resize', action='store_true', default=False,
+                        help='Resize the images to a given dimension.')
+    parser.add_argument('--h', type=int, default='480',
+                        help='Image height.')
+    parser.add_argument('--w', type=int, default='640',
+                        help='Image width.')
+    parser.add_argument('--extension', type=str, default=None,
+                         help="Extension to add to each exported npz.")
+    args = parser.parse_args()
+
+    model = args.model
+    if model not in ['lisrd', 'lisrd_sift']:
+        raise ValueError(f"Unknown model: {model}")
+    checkpoint = os.path.expanduser(args.checkpoint)
+    if not os.path.exists(checkpoint):
+        sys.exit('Unable to find checkpoint', checkpoint)
+
+    keypoints_type = args.keypoints
+    if keypoints_type not in ['sift', 'superpoint']:
+        sys.exit('Unknown keypoint method:', keypoints_type)
+    num_keypoints = args.num_kp
+
+    extension = args.extension if args.extension else model
+    extension = '.' + extension
+
+    export(args.images_list, model, checkpoint, keypoints_type,
+           num_keypoints, args.detection_thresh, extension,
+           args.resize, args.h, args.w)

vismatch/third_party/LISRD/lisrd/models/__init__.py

@@ -0,0 +1,7 @@
+def get_model(name):
+    mod = __import__('lisrd.models.{}'.format(name), fromlist=[''])
+    return getattr(mod, _module_to_class(name))
+
+
+def _module_to_class(name):
+    return ''.join(n.capitalize() for n in name.split('_'))

vismatch/third_party/LISRD/lisrd/models/backbones/net_vlad.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as func
+
+
+class NetVLAD(nn.Module):
+    """
+    NetVLAD layer implementation
+    Credits: https://github.com/lyakaap/NetVLAD-pytorch
+    """
+
+    def __init__(self, num_clusters=64, dim=128, alpha=100.0,
+                 normalize_input=True):
+        """
+        Args:
+            num_clusters: number of clusters.
+            dim: dimension of descriptors.
+            alpha: parameter of initialization. Larger is harder assignment.
+            normalize_input: if true, descriptor-wise L2 normalization
+                             is applied to input.
+        """
+        super().__init__()
+        self.num_clusters = num_clusters
+        self.dim = dim
+        self.alpha = alpha
+        self.normalize_input = normalize_input
+        self.conv = nn.Conv2d(dim, num_clusters,
+                              kernel_size=(1, 1), bias=True)
+        self.centroids = nn.Parameter(torch.rand(num_clusters, dim))
+        self._init_params()
+
+    def _init_params(self):
+        self.conv.weight = nn.Parameter(
+            (2.0 * self.alpha * self.centroids).unsqueeze(-1).unsqueeze(-1))
+        self.conv.bias = nn.Parameter(
+            - self.alpha * self.centroids.norm(dim=1))
+
+    def forward(self, x):
+        N, C = x.shape[:2]
+
+        if self.normalize_input:
+            x = func.normalize(x, p=2, dim=1)  # across descriptor dim
+
+        # Soft-assignment
+        soft_assign = self.conv(x).view(N, self.num_clusters, -1)
+        soft_assign = func.softmax(soft_assign, dim=1)
+
+        x_flatten = x.view(N, C, -1)
+        
+        # Calculate residuals to each clusters
+        residual = (x_flatten.expand(
+            self.num_clusters, -1, -1, -1).permute(1, 0, 2, 3)
+            - self.centroids.expand(
+                x_flatten.size(-1), -1, -1).permute(1, 2, 0).unsqueeze(0))
+        residual *= soft_assign.unsqueeze(2)
+        vlad = residual.sum(dim=-1)
+
+        vlad = func.normalize(vlad, p=2, dim=2)  # intra-normalization
+        vlad = vlad.view(x.size(0), -1)  # flatten
+        vlad = func.normalize(vlad, p=2, dim=1)  # L2 normalize
+
+        return vlad

vismatch/third_party/LISRD/lisrd/models/backbones/vgg.py

@@ -0,0 +1,46 @@
+import torch
+
+
+class VGGLikeModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._relu = torch.nn.ReLU(inplace=True)
+        self._pool = torch.nn.AvgPool2d(kernel_size=2, stride=2)
+        self._conv1_1 = torch.nn.Conv2d(3, 64, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn1_1 = torch.nn.BatchNorm2d(64)
+        self._conv1_2 = torch.nn.Conv2d(64, 64, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn1_2 = torch.nn.BatchNorm2d(64)
+        self._conv2_1 = torch.nn.Conv2d(64, 64, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn2_1 = torch.nn.BatchNorm2d(64)
+        self._conv2_2 = torch.nn.Conv2d(64, 64, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn2_2 = torch.nn.BatchNorm2d(64)
+        self._conv3_1 = torch.nn.Conv2d(64, 128, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn3_1 = torch.nn.BatchNorm2d(128)
+        self._conv3_2 = torch.nn.Conv2d(128, 128, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn3_2 = torch.nn.BatchNorm2d(128)
+        self._conv4_1 = torch.nn.Conv2d(128, 256, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn4_1 = torch.nn.BatchNorm2d(256)
+        self._conv4_2 = torch.nn.Conv2d(256, 256, kernel_size=3,
+                                        stride=1, padding=1)
+        self._bn4_2 = torch.nn.BatchNorm2d(256)
+
+    def forward(self, inputs):
+        x = self._bn1_1(self._relu(self._conv1_1(inputs)))
+        x = self._bn1_2(self._relu(self._conv1_2(x)))
+        x = self._pool(x)
+        x = self._bn2_1(self._relu(self._conv2_1(x)))
+        x = self._bn2_2(self._relu(self._conv2_2(x)))
+        x = self._pool(x)
+        x = self._bn3_1(self._relu(self._conv3_1(x)))
+        x = self._bn3_2(self._relu(self._conv3_2(x)))
+        x = self._pool(x)
+        x = self._bn4_1(self._relu(self._conv4_1(x)))
+        x = self._bn4_2(self._relu(self._conv4_2(x)))
+        return x

vismatch/third_party/LISRD/lisrd/models/base_model.py

@@ -0,0 +1,336 @@
+import warnings
+warnings.filterwarnings(action='once')
+warnings.filterwarnings("ignore", message="numpy.ufunc size changed")
+import os
+import logging
+import numpy as np
+import torch
+import torch.optim as optim
+from tqdm import tqdm
+from abc import ABCMeta, abstractmethod
+from torch.utils.tensorboard import SummaryWriter
+
+logging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
+
+
+class Mode:
+    TRAIN = 'train'
+    VAL = 'validation'
+    TEST = 'test'
+    EXPORT = 'export'
+
+
+class BaseModel(metaclass=ABCMeta):
+    """Base model class.
+
+    Arguments:
+        dataset: A BaseDataset object.
+        config: A dictionary containing the configuration parameters.
+                The Entry `learning_rate` is required.
+
+    Models should inherit from this class and implement the following methods:
+        `_model`, `_loss`, `_metrics` and `initialize_weights`.
+
+    Additionally, the static attribute required_config_keys is a list
+    containing the required config entries.
+    """
+    required_baseconfig = ['learning_rate']
+
+    @abstractmethod
+    def _model(self, config):
+        """ Implements the model.
+
+        Arguments:
+            config: A configuration dictionary.
+
+        Returns:
+            A torch.nn.Module that implements the model.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _forward(self, inputs, mode, config):
+        """ Calls the model on some input.
+        This method is called three times: for training, testing and
+        prediction (see the `mode` argument) and can return different tensors
+        depending on the mode. It only supports NCHW format for now.
+
+        Arguments:
+            inputs: A dictionary of input features, where the keys are their
+                names (e.g. `"image"`) and the values of type `torch.Tensor`.
+            mode: An attribute of the `Mode` class, either `Mode.TRAIN`,
+                  `Mode.TEST` or `Mode.PRED`.
+            config: A configuration dictionary.
+
+        Returns:
+            A dictionary of outputs, where the keys are their names
+            (e.g. `"logits"`) and the values are the corresponding Tensor.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _loss(self, outputs, inputs, config):
+        """ Implements the training loss.
+        This method is called on the outputs of the `_model` method
+        in training mode.
+
+        Arguments:
+            outputs: A dictionary, as returned by `_model` called with
+                     `mode=Mode.TRAIN`.
+            inputs: A dictionary of input features (same as for `_model`).
+            config: A configuration dictionary.
+
+        Returns:
+            A Tensor corresponding to the loss to minimize during training.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _metrics(self, outputs, inputs, config):
+        """ Implements the evaluation metrics.
+        This method is called on the outputs of the `_model` method
+        in test mode.
+
+        Arguments:
+            outputs: A dictionary, as returned by `_model` called with
+                     `mode=Mode.EVAL`.
+            inputs: A dictionary of input features (same as for `_model`).
+            config: A configuration dictionary.
+
+        Returns:
+            A dictionary of metrics, where the keys are their names
+            (e.g. "`accuracy`") and the values are the corresponding Tensor.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def initialize_weights(self):
+        """ Initialize all the weights in the network. """
+        return NotImplementedError
+
+    def __init__(self, dataset, config, device):
+        self._dataset = dataset
+        self._config = config
+        required = self.required_baseconfig + getattr(
+            self, 'required_config_keys', [])
+        for r in required:
+            assert r in self._config, 'Required configuration entry: \'{}\''.format(r)
+        self._net = self._model(config)
+        if torch.cuda.device_count() > 1:
+            logging.info('Using {} GPU(s).'.format(torch.cuda.device_count()))
+            self._net = torch.nn.DataParallel(self._net)
+        self._net = self._net.to(device)
+        self._solver = optim.Adam(self._net.parameters(),
+                                  lr=self._config['learning_rate'])
+        self._it = 0
+        self._epoch = 0
+
+    def _to_dict(self, d, device):
+        """ Send all the values of a dict of Tensor to a specific device. """
+        return {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                for k, v in d.items()}
+
+    def train(self, n_iter, n_epoch, exper_dir, validation_interval=100,
+              save_interval=500, resume_training='', device='cpu'):
+        """ Train the model.
+
+        Arguments:
+            n_iter: max number of iterations.
+            n_epoch: max number of epochs.
+            exper_dir: folder containing the outputs of the training.
+            validation_interval: evaluate the model every
+                                 'validation_interval' iter.
+            save_interval: save the model every 'save_interval' iter.
+            resume_training: path to the checkpoint file when training is
+                             resumed from a previous training.
+            device: device on which to perform the operations.
+        """
+        self.initialize_weights()
+        if resume_training == '':
+            logging.info('Initializing new weights.')
+            self._it = 0
+            self._epoch = 0
+        else:
+            logging.info('Loading weights from ' + resume_training)
+            self.load(resume_training)
+        self._net.train()
+
+        runs_dir = os.path.join(exper_dir, 'runs')
+        self._writer = SummaryWriter(runs_dir)
+        train_data_loader = self._dataset.get_data_loader('train')
+        val_data_loader = self._dataset.get_data_loader('val')
+
+        train_loss = []
+        logging.info('Start training')
+        while self._epoch < n_epoch:
+            for x in train_data_loader:
+                if self._it > n_iter:
+                    break
+
+                # Optimize one batch
+                inputs = self._to_dict(x, device)
+                outputs = self._forward(inputs, Mode.TRAIN, self._config)
+                loss = self._loss(outputs, inputs, self._config)
+                train_loss.append(loss.item())
+                self._solver.zero_grad()
+                loss.backward()
+                self._solver.step()
+
+                # Validation
+                if self._it % validation_interval == 0:
+                    validation_loss = []
+                    metrics = {}
+                    for x in val_data_loader:
+                        inputs = self._to_dict(x, device)
+                        outputs = self._forward(inputs, Mode.VAL,
+                                                self._config)
+                        loss = self._loss(outputs, inputs, self._config)
+                        validation_loss.append(loss.item())
+                        metric = self._metrics(outputs, inputs, self._config)
+                        metrics = {k: [metric[k].item()] if k not in metrics
+                                   else metrics[k] + [metric[k].item()]
+                                   for k in metric}
+                    total_train_loss = np.mean(train_loss)
+                    total_validation_loss = np.mean(validation_loss)
+                    total_metrics = {k: np.mean(v)
+                                     for k, v in metrics.items()}
+
+                    # Print the metrics
+                    logging.info(
+                        'Iter {:4d}: train loss {:.4f}, validation loss {:.4f}'.format(
+                            self._it, total_train_loss, total_validation_loss)
+                        + ''.join([', {} {:.4f}'.format(m, total_metrics[m])
+                                   for m in total_metrics]))
+
+                    # Save them a summary files
+                    self._writer.add_scalar(
+                        'train_loss', total_train_loss, self._it)
+                    self._writer.add_scalar(
+                        'validation_loss', total_validation_loss, self._it)
+                    for m in total_metrics:
+                        self._writer.add_scalar(m, total_metrics[m], self._it)
+                
+                if self._it % save_interval == 0:
+                    self.save(exper_dir)
+
+                self._it += 1
+            self._epoch += 1
+            if self._it > n_iter:
+                break
+        self._writer.close()
+        logging.info('Training finished.')
+    
+    def test(self, exper_dir, checkpoint_path, device='cpu'):
+        """ Test the model on a test dataset.
+
+        Arguments:
+            exper_dir: folder containing the outputs of the training.
+            checkpoint_path: path to the checkpoint.
+            device: device on which to perform the operations.
+        """
+        # Load the weights
+        logging.info('Loading weights from ' + checkpoint_path)
+        self.load(checkpoint_path, Mode.TEST)
+        self._net.eval()
+        test_data_loader = self._dataset.get_data_loader('test')
+
+        # Run the evaluation metrics on the test dataset
+        logging.info('Start evaluation.')
+        metrics = {}
+        for x in tqdm(test_data_loader):
+            inputs = self._to_dict(x, device)
+            outputs = self._forward(inputs, Mode.TEST,
+                                    self._config)
+            metric = self._metrics(outputs, inputs, self._config)
+            metrics = {k: [metric[k].item()] if k not in metrics
+                        else metrics[k] + [metric[k].item()]
+                        for k in metric}
+        total_metrics = {k: np.mean(v) for k, v in metrics.items()}
+        logging.info('Test metrics: '
+                     + ''.join(['{} {:.4f};'.format(m, total_metrics[m])
+                                for m in total_metrics]))
+    
+    def export(self, exper_dir, checkpoint_path, output_dir, device='cpu'):
+            """ Export the descriptors on a given dataset.
+
+            Arguments:
+                exper_dir: folder containing the outputs of the training.
+                checkpoint_path: path to the checkpoint.
+                output_dir: for each item in the dataset, write a .npz file
+                in output_dir with the exported descriptors.
+                device: device on which to perform the operations.
+            """
+            # Load the weights
+            logging.info('Loading weights from ' + checkpoint_path)
+            self.load(checkpoint_path, Mode.TEST)
+            self._net.eval()
+            test_data_loader = self._dataset.get_data_loader('test')
+
+            # Run the evaluation metrics on the test dataset
+            logging.info('Start exporting.')
+            i = 0
+            for x in tqdm(test_data_loader):
+                inputs = self._to_dict(x, device)
+                outputs = self._forward(inputs, Mode.EXPORT,
+                                        self._config)
+                outputs.update(inputs)
+                for k, v in outputs.items():
+                    outputs[k] = v.detach().cpu().numpy()[0]
+                    if len(outputs[k].shape) == 3:
+                        outputs[k] = outputs[k].transpose(1, 2, 0)
+                out_file = os.path.join(output_dir, str(i) + '.npz')
+                np.savez_compressed(out_file, **outputs)
+                i += 1
+            logging.info('Descriptors exported in ' + output_dir)
+
+    def _adapt_weight_names(self, state_dict):
+        """ Adapt the weight names when the training and testing are done
+        with a different GPU configuration (with/without DataParallel). """
+        train_parallel = list(state_dict.keys())[0][:7] == 'module.'
+        test_parallel = torch.cuda.device_count() > 1
+        new_state_dict = {}
+        if train_parallel and (not test_parallel):
+            # Need to remove 'module.' from all the variable names
+            for k, v in state_dict.items():
+                new_state_dict[k[7:]] = v
+        elif test_parallel and (not train_parallel):
+            # Need to add 'module.' to all the variable names
+            for k, v in state_dict.items():
+                new_k = 'module.' + k
+                new_state_dict[new_k] = v
+        else:  # Nothing to do
+            new_state_dict = state_dict
+        return new_state_dict
+
+    def _match_state_dict(self, old_state_dict, new_state_dict):
+        """ Return a new state dict that has exactly the same entries
+        as old_state_dict and that is updated with the values of
+        new_state_dict whose entries are shared with old_state_dict.
+        This allows loading a pre-trained network. """
+        return ({k: new_state_dict[k] if k in new_state_dict else v
+                 for (k, v) in old_state_dict.items()},
+                old_state_dict.keys() == new_state_dict.keys())
+
+    def save(self, exper_dir):
+        """ Save the current training in a .pth file. """
+        save_file = os.path.join(
+            exper_dir, 'checkpoints/checkpoint_' + str(self._it) + '.pth')
+        torch.save({'iter': self._it,
+                    'model_state_dict': self._net.state_dict(),
+                    'optimizer_state_dict': self._solver.state_dict()},
+                   save_file)
+
+    def load(self, checkpoint_path, mode=Mode.TRAIN):
+        """ Load a model stored on disk. """
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')
+        updated_state_dict, same_net = self._match_state_dict(
+            self._net.state_dict(),
+            self._adapt_weight_names(checkpoint['model_state_dict']))
+        self._net.load_state_dict(updated_state_dict)
+        if same_net:
+            self._solver.load_state_dict(checkpoint['optimizer_state_dict'])
+        self._it = checkpoint['iter']
+        if mode == Mode.TRAIN:
+            self._epoch = (self._it * self._dataset._config['batch_size']
+                           // len(self._dataset))

vismatch/third_party/LISRD/lisrd/models/keypoint_detectors.py

@@ -0,0 +1,34 @@
+import os
+import numpy as np
+import cv2
+import torch
+
+from ..third_party.super_point_magic_leap.demo_superpoint import SuperPointFrontend
+
+
+def SIFT_detect(img, nfeatures=1500, contrastThreshold=0.04):
+    """ Compute SIFT feature points. """
+    sift = cv2.SIFT_create(nfeatures=nfeatures,
+                                       contrastThreshold=contrastThreshold)
+    keypoints = sift.detect(img, None)
+    keypoints = [[k.pt[1], k.pt[0], k.response] for k in keypoints]
+    keypoints = np.array(keypoints)
+    return keypoints
+
+
+def SP_detect(img, kp_net):
+    """ Compute SuperPoint feature points. """
+    keypoints, _, _ = kp_net.run(img.astype(np.float32) / 255.)
+    keypoints = keypoints.squeeze()[[1, 0, 2], :].transpose()
+    return keypoints
+
+
+def load_SP_net(conf_thresh=0.015, cuda=torch.cuda.is_available(),
+                nms_dist=4, nn_thresh=0.7):
+    weights_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        '../third_party/super_point_magic_leap/superpoint_v1.pth')
+    kp_net = SuperPointFrontend(
+        weights_path, nms_dist=nms_dist, conf_thresh=conf_thresh,
+        nn_thresh=nn_thresh, cuda=cuda)
+    return kp_net