load from github release

Author: leoluxxx

scripts/reproduce_test/outdoor.sh

@@ -9,7 +9,7 @@ cd $PROJECT_DIR
 
 data_cfg_path="configs/data/megadepth_test_1500.py"
 main_cfg_path="configs/jamma/outdoor/test.py"
-ckpt_path='weight/jamma_weight.ckpt'
+ckpt_path='official' # your path or 'official' (load from github release)
 dump_dir="dump/jamma_outdoor"
 profiler_name="inference"
 n_nodes=1  # mannually keep this the same with --nodes

src/jamma/backbone.py

@@ -16,6 +16,11 @@ def __init__(self):
         self.cnn.stages[2] = None
         self.cnn.stages[3] = None
 
+        state_dict = torch.hub.load_state_dict_from_url(
+            'https://github.com/leoluxxx/JamMa/releases/download/v0.1/convnextv2_nano_pretrain.ckpt',
+            file_name='convnextv2_nano_pretrain.ckpt')
+        self.cnn.load_state_dict(state_dict, strict=True)
+
         self.lin_4 = nn.Conv2d(80, 128, 1)
         self.lin_8 = nn.Conv2d(160, 256, 1)
 

src/lightning/lightning_jamma.py

@@ -39,8 +39,14 @@ def __init__(self, config, pretrained_ckpt=None, profiler=None, dump_dir=None):
         self.matcher = JamMa(config=_config['jamma'], profiler=profiler)
         self.loss = Loss(_config)
 
-        if pretrained_ckpt:
-            state_dict = torch.load(pretrained_ckpt, map_location='cpu')
+        if pretrained_ckpt == 'official':
+            state_dict = torch.hub.load_state_dict_from_url(
+                'https://github.com/leoluxxx/JamMa/releases/download/v0.1/jamma.ckpt',
+                file_name='jamma.ckpt')['state_dict']
+            self.load_state_dict(state_dict, strict=True)
+            logger.info(f"Load Official JamMa Weight")
+        elif pretrained_ckpt:
+            state_dict = torch.load(pretrained_ckpt, map_location='cpu')['state_dict']
             self.load_state_dict(state_dict, strict=True)
             logger.info(f"Load \'{pretrained_ckpt}\' as pretrained checkpoint")