Update dependencies and improve documentation (#34)

* Update dependencies and improve documentation

- Updated flake8 to version 5.0.4 in .pre-commit-config.yaml.
- Added gdown to requirements.txt for model downloading.
- Changed dataset download link in get_started.md and conftest.py to use Google Drive.
- Refactored model download logic in bmnet.py to use gdown.
- Introduced gdown_download_model function in miscs.py for consistent model downloading.

* Update lint workflow to use Ubuntu 22.04 and remove unnecessary Ruby installation steps

Author: LazyBusyYang

.github/workflows/lint.yml

@@ -8,7 +8,7 @@ concurrency:
 
 jobs:
   lint:
-    runs-on: ubuntu-18.04
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python 3.7
@@ -17,9 +17,6 @@ jobs:
           python-version: 3.7
       - name: Install pre-commit hook
         run: |
-          sudo apt-add-repository ppa:brightbox/ruby-ng -y
-          sudo apt-get update
-          sudo apt-get install -y ruby2.7
           pip install pre-commit
           pre-commit install
       - name: Linting

.pre-commit-config.yaml

@@ -4,8 +4,8 @@ repos:
     rev: v14.0.1
     hooks:
       - id: clang-format
-  - repo: https://gitlab.com/pycqa/flake8.git
-    rev: 3.8.3
+  - repo: https://github.com/PyCQA/flake8
+    rev: 5.0.4
     hooks:
       - id: flake8
   - repo: https://github.com/asottile/seed-isort-config.git

docs/en/get_started.md

@@ -2,7 +2,7 @@
 XRLocalization provides a flexible tool that can easily perform visual localization offline and online.
 Given a query image, XRLocalization estimates a 6DoF pose from a pre-reconstructed map. A tiny dataset
 is provided for convenience. Download the dataset from
-[here](https://openxrlab-share-mainland.oss-cn-hangzhou.aliyuncs.com/xrlocalization/meta/xrloc-test-meta.tar.gz). The
+[here](https://drive.google.com/file/d/1vKfCDWtZ1ui5t5sYjlF_EAqj1mVM1zk-/view?usp=sharing). The
 folder extracted from the dataset is shown below.
 ```commandline
 ├── map

requirements.txt

@@ -9,4 +9,5 @@ scipy>=1.6.1
 setuptools>=46.1.3
 torch>=1.1.0
 tqdm
-xrprimer==0.5.2
+xrprimer==0.5.2
+gdown

tests/conftest.py

@@ -10,8 +10,8 @@
 def fixture():
     if os.path.exists(dataset_dir):
         shutil.rmtree(dataset_dir)
-    url = 'https://openxrlab-share-mainland.oss-cn-hangzhou.aliyuncs.com/xrlocalization/meta/xrloc-test-meta.tar.gz'
-    command = ['wget', '--no-check-certificate', url]
+    url = 'https://docs.google.com/uc?id=1vKfCDWtZ1ui5t5sYjlF_EAqj1mVM1zk-'
+    command = ['gdown', url]
     subprocess.run(command, check=True)
     command = ['tar', '-xf', 'xrloc-test-meta.tar.gz']
     subprocess.run(command, check=True)

xrloc/localizer.py

@@ -37,21 +37,23 @@ def __init__(self, map_path, config=default_config):
         head_logging('XRLocalization')
         self.config = config
         database_path = os.path.join(map_path, 'database.bin')
-        
+
         if os.path.exists(database_path):
             self.database = ImageDatabase(database_path)
             self.database.create()
         else:
-            pairs = [name for name in os.listdir(map_path) 
-                     if name.startswith('pairs-query')]
+            pairs = [
+                name for name in os.listdir(map_path)
+                if name.startswith('pairs-query')
+            ]
             if len(pairs) == 0:
                 raise ValueError(
                     'Not found database under map: {}'.format(map_path))
             else:
                 self.pairs = PairsDatabase(os.path.join(map_path, pairs[0]))
 
         self.reconstruction = Reconstruction(map_path)
-        
+
         if hasattr(self, 'database'):
             self.gextractor = Extractor(self.config['global_feature'])
         self.lextractor = Extractor(self.config['local_feature'])
@@ -61,8 +63,9 @@ def __init__(self, map_path, config=default_config):
         config_logging(self.config)
 
         if self.config['mode'] == '2D2D' and self.config['matcher'] == 'gam':
-            raise ValueError('Loc mode {} is not compatible with matcher {}'.format(
-                self.config['mode'], self.config['matcher']))
+            raise ValueError(
+                'Loc mode {} is not compatible with matcher {}'.format(
+                    self.config['mode'], self.config['matcher']))
 
         head_logging('Init Success')
 
@@ -91,15 +94,19 @@ def geo_localize(self, data):
         if hasattr(self, 'database'):
             image_feature = self.gextractor.extract(data)
             image_ids = self.database.retrieve(image_feature,
-                                            self.config['retrieval_num'])
+                                               self.config['retrieval_num'])
             return image_ids
         elif hasattr(self, 'pairs') and isinstance(data, str):
-            image_names = self.pairs.image_retrieve(data, self.config['retrieval_num'])
-            image_ids = [self.reconstruction.name_to_id(name) for name in image_names]
+            image_names = self.pairs.image_retrieve(
+                data, self.config['retrieval_num'])
+            image_ids = [
+                self.reconstruction.name_to_id(name) for name in image_names
+            ]
             return np.array([id for id in image_ids if id != -1])
 
     def feature_match_2d3d(self, query_points, query_point_descriptors,
-                           train_points, train_point_descriptors, width, height):
+                           train_points, train_point_descriptors, width,
+                           height):
         """Feature matching phase."""
         query_feat = {
             'shape': np.array([height, width]),
@@ -114,58 +121,57 @@ def feature_match_2d3d(self, query_points, query_point_descriptors,
         return pred['matches'], pred['scores']
 
     def establish_correspondences_2d2d(self, query_feat, image_ids):
-        '''Establish 2D-3D correspondences depend on 2D2D matching
-        '''
+        """Establish 2D-3D correspondences depend on 2D2D matching."""
         logging.info('Scene size: {0}'.format(len(image_ids)))
-        match_indices = np.ones(len(query_feat['points']), dtype=int)*-1
+        match_indices = np.ones(len(query_feat['points']), dtype=int) * -1
         match_priors = np.zeros(len(query_feat['points']))
         for image_id in image_ids:
             ref_image = self.reconstruction.image_at(image_id)
             ref_feat = {
-                'points': ref_image.xys,
-                'descs': self.reconstruction.point3d_features(ref_image.point3D_ids),
-                'scores': np.ones(len(ref_image.xys)),
-                'shape': np.array([600, 600]) # TODO
+                'points':
+                ref_image.xys,
+                'descs':
+                self.reconstruction.point3d_features(ref_image.point3D_ids),
+                'scores':
+                np.ones(len(ref_image.xys)),
+                'shape':
+                np.array([600, 600])  # TODO
             }
             pred = self.matcher.match(query_feat, ref_feat)
             matches, scores = pred['matches'], pred['scores']
-            
+
             reserve_matches = matches[:, scores > match_priors[matches[0]]]
             reserve_scores = scores[scores > match_priors[matches[0]]]
             if len(reserve_scores) > 0:
-                match_indices[reserve_matches[0]] = ref_image.point3D_ids[reserve_matches[1]]
+                match_indices[reserve_matches[0]] = ref_image.point3D_ids[
+                    reserve_matches[1]]
                 match_priors[reserve_matches[0]] = reserve_scores
 
             if len(match_priors[match_indices != -1]) > 400:
                 break
 
         point3d_ids = match_indices[match_indices != -1]
-        points3d = self.reconstruction.point3d_coordinates(
-            point3d_ids)
+        points3d = self.reconstruction.point3d_coordinates(point3d_ids)
         points2d = query_feat['points'][match_indices != -1]
         priors = match_priors[match_indices != -1]
         logging.info('Match number: {0}'.format(len(priors)))
         return points2d, points3d, priors
 
-
     def establish_correspondences_2d3d(self, feat, image_ids):
-        '''Establish 2D-3D correspondences depend on 2D3D matching
-        '''
+        """Establish 2D-3D correspondences depend on 2D3D matching."""
         point3d_ids = self.reconstruction.visible_points(image_ids)
-        point3ds = self.reconstruction.point3d_coordinates(
-            point3d_ids)
-        point3d_descs = self.reconstruction.point3d_features(
-            point3d_ids)
-        logging.info('3d points size: {0}'.format(
-            point3d_descs.shape[1]))
+        point3ds = self.reconstruction.point3d_coordinates(point3d_ids)
+        point3d_descs = self.reconstruction.point3d_features(point3d_ids)
+        logging.info('3d points size: {0}'.format(point3d_descs.shape[1]))
 
         # Matching
-        matches, priors = self.feature_match_2d3d(feat['points'],
-                                                  feat['descs'],
-                                                  point3ds,
-                                                  point3d_descs,
-                                                  feat['shape'][1], # width
-                                                  feat['shape'][0])
+        matches, priors = self.feature_match_2d3d(
+            feat['points'],
+            feat['descs'],
+            point3ds,
+            point3d_descs,
+            feat['shape'][1],  # width
+            feat['shape'][0])
         logging.info('Match number: {0}'.format(matches.shape[1]))
         points2d = feat['points'][matches[0]]
         points3d = point3ds[matches[1]]
@@ -202,19 +208,22 @@ def refine_localize(self, image, camera, ref_image_ids):
         logging.info('Coarse location number: {0}'.format(len(scenes)))
 
         best_ret = {
-            'ninlier': 0, 'qvec': np.array([1, 0, 0, 0]),
-            'tvec': np.array([0, 0, 0]), 'mask': None
+            'ninlier': 0,
+            'qvec': np.array([1, 0, 0, 0]),
+            'tvec': np.array([0, 0, 0]),
+            'mask': None
         }
         for i, image_ids in enumerate(scenes[:self.config['max_scene_num']]):
             # Establish 2D-3D correspondences
             if self.config['mode'] == '2D3D':
-                points2d, points3d, priors = self.establish_correspondences_2d3d(
-                    feat, image_ids)
+                points2d, points3d, priors = \
+                    self.establish_correspondences_2d3d(feat, image_ids)
             elif self.config['mode'] == '2D2D':
-                points2d, points3d, priors = self.establish_correspondences_2d2d(
-                    feat, image_ids)
+                points2d, points3d, priors = \
+                    self.establish_correspondences_2d2d(feat, image_ids)
 
-            if len(priors) < 3: continue
+            if len(priors) < 3:
+                continue
 
             # Pose estimation
             ret = self.prior_guided_pose_estimation(points2d, points3d, priors,
@@ -226,4 +235,4 @@ def refine_localize(self, image, camera, ref_image_ids):
                 best_ret = ret
             if best_ret['ninlier'] > self.config['max_inlier']:
                 break
-        return best_ret
+        return best_ret

xrloc/matchers/bmnet.py

@@ -3,7 +3,7 @@
 
 import torch.nn as nn
 
-from xrloc.utils.miscs import download_model, get_parent_dir
+from xrloc.utils.miscs import gdown_download_model, get_parent_dir
 from scipy.optimize import linear_sum_assignment
 
 
@@ -135,11 +135,11 @@ def __init__(self, config=default_config):
         # 'hpooling'] else None
         self.conv = nn.Conv1d(self.config['channels2'][-1], 1, kernel_size=1)
 
-        url = 'https://openxrlab-share-mainland.oss-cn-hangzhou.aliyuncs.com/xrlocalization/weights/bmnet.pth'
+        url = 'https://docs.google.com/uc?id=1RqkkGc5WomkP7aDgLbCXLqlc8jcwdlcN'
         model_dir = get_parent_dir(__file__) + '/../models/'
 
         model_name = 'bmnet.pth'
-        download_model(url, model_dir, model_name)
+        gdown_download_model(url, model_dir, model_name)
         model_path = os.path.join(model_dir, model_name)
 
         self.load_state_dict(

xrloc/utils/miscs.py

@@ -107,6 +107,16 @@ def download_model(url, model_dir, model_name):
     subprocess.run(command, check=True)
 
 
+def gdown_download_model(url, model_dir, model_name):
+    model_path = os.path.join(model_dir, model_name)
+    if os.path.exists(model_path):
+        return
+    os.makedirs(model_dir, exist_ok=True)
+    print('Downloading the {} model from {}.'.format(model_name, url))
+    command = ['gdown', url, '-O', model_path]
+    subprocess.run(command, check=True)
+
+
 def head_logging(info: str, width=50):
     logging.info('=' * width)
     left = '*' * int((width - len(info) - 2) / 2) + ' '