rename: roma -> romatch (#50)

README.md

@@ -34,7 +34,7 @@ pip install -e .
 We provide two demos in the [demos folder](demo).
 Here's the gist of it:
 ```python
-from roma import roma_outdoor
+from romatch import roma_outdoor
 roma_model = roma_outdoor(device=device)
 # Match
 warp, certainty = roma_model.match(imA_path, imB_path, device=device)
@@ -48,7 +48,7 @@ F, mask = cv2.findFundamentalMat(
 )
 ```
 
-**New**: You can also match arbitrary keypoints with RoMa. See [match_keypoints](roma/models/matcher.py) in RegressionMatcher.
+**New**: You can also match arbitrary keypoints with RoMa. See [match_keypoints](romatch/models/matcher.py) in RegressionMatcher.
 
 ## Settings
 
@@ -84,7 +84,7 @@ Our codebase builds on the code in [DKM](https://github.com/Parskatt/DKM).
 ## Tiny RoMa
 If you find that RoMa is too heavy, you might want to try Tiny RoMa which is built on top of XFeat.
 ```python
-from roma import tiny_roma_v1_outdoor
+from romatch import tiny_roma_v1_outdoor
 tiny_roma_model = tiny_roma_v1_outdoor(device=device)
 ```
 Mega1500:

demo/demo_3D_effect.py

@@ -2,9 +2,9 @@ from PIL import Image
 import torch
 import torch.nn.functional as F
 import numpy as np
-from roma.utils.utils import tensor_to_pil
+from romatch.utils.utils import tensor_to_pil
 
-from roma import roma_outdoor
+from romatch import roma_outdoor
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 

demo/demo_fundamental.py

@@ -1,7 +1,7 @@
 from PIL import Image
 import torch
 import cv2
-from roma import roma_outdoor
+from romatch import roma_outdoor
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 

demo/demo_match.py

@@ -2,9 +2,9 @@ from PIL import Image
 import torch
 import torch.nn.functional as F
 import numpy as np
-from roma.utils.utils import tensor_to_pil
+from romatch.utils.utils import tensor_to_pil
 
-from roma import roma_outdoor
+from romatch import roma_outdoor
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 

experiments/eval_roma_outdoor.py

@@ -9,9 +9,9 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 import json
 import wandb
 
-from roma.benchmarks import MegadepthDenseBenchmark
-from roma.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
-from roma.benchmarks import Mega1500PoseLibBenchmark
+from romatch.benchmarks import MegadepthDenseBenchmark
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
+from romatch.benchmarks import Mega1500PoseLibBenchmark
 
 def test_mega_8_scenes(model, name):
     mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
@@ -57,7 +57,7 @@ def test_hpatches(model, name):
 
 
 if __name__ == "__main__":
-    from roma import roma_outdoor
+    from romatch import roma_outdoor
     model = roma_outdoor(device = "cuda", coarse_res = 672, upsample_res = 1344)
     experiment_name = "roma_latest"
     #test_mega1500(model, experiment_name)

experiments/eval_tiny_roma_v1_outdoor.py

@@ -13,9 +13,9 @@ from torch.utils.data import ConcatDataset
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 import json
-from roma.benchmarks import ScanNetBenchmark
-from roma.benchmarks import Mega1500PoseLibBenchmark, ScanNetPoselibBenchmark
-from roma.benchmarks import MegaDepthPoseEstimationBenchmark
+from romatch.benchmarks import ScanNetBenchmark
+from romatch.benchmarks import Mega1500PoseLibBenchmark, ScanNetPoselibBenchmark
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark
 
 def test_mega_8_scenes(model, name):
     mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
@@ -85,7 +85,7 @@ if __name__ == "__main__":
     os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
     os.environ["OMP_NUM_THREADS"] = "16"
     torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
-    from roma import tiny_roma_v1_outdoor
+    from romatch import tiny_roma_v1_outdoor
 
     experiment_name = Path(__file__).stem
     device = 'cuda'

experiments/roma_indoor.py

@@ -11,16 +11,16 @@ import json
 import wandb
 from tqdm import tqdm
 
-from roma.benchmarks import MegadepthDenseBenchmark
-from roma.datasets.megadepth import MegadepthBuilder
-from roma.datasets.scannet import ScanNetBuilder
-from roma.losses.robust_loss import RobustLosses
-from roma.benchmarks import MegadepthDenseBenchmark, ScanNetBenchmark
-from roma.train.train import train_k_steps
-from roma.models.matcher import *
-from roma.models.transformer import Block, TransformerDecoder, MemEffAttention
-from roma.models.encoders import *
-from roma.checkpointing import CheckPoint
+from romatch.benchmarks import MegadepthDenseBenchmark
+from romatch.datasets.megadepth import MegadepthBuilder
+from romatch.datasets.scannet import ScanNetBuilder
+from romatch.losses.robust_loss import RobustLosses
+from romatch.benchmarks import MegadepthDenseBenchmark, ScanNetBenchmark
+from romatch.train.train import train_k_steps
+from romatch.models.matcher import *
+from romatch.models.transformer import Block, TransformerDecoder, MemEffAttention
+from romatch.models.encoders import *
+from romatch.checkpointing import CheckPoint
 
 resolutions = {"low":(448, 448), "medium":(14*8*5, 14*8*5), "high":(14*8*6, 14*8*6)}
 
@@ -174,14 +174,14 @@ def train(args):
     rank = dist.get_rank()
     print(f"Start running DDP on rank {rank}")
     device_id = rank % torch.cuda.device_count()
-    roma.LOCAL_RANK = device_id
+    romatch.LOCAL_RANK = device_id
     torch.cuda.set_device(device_id)
     
     resolution = args.train_resolution
     wandb_log = not args.dont_log_wandb
     experiment_name = os.path.splitext(os.path.basename(__file__))[0]
     wandb_mode = "online" if wandb_log and rank == 0 and False else "disabled"
-    wandb.init(project="roma", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
+    wandb.init(project="romatch", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
     checkpoint_dir = "workspace/checkpoints/"
     h,w = resolutions[resolution]
     model = get_model(pretrained_backbone=True, resolution=resolution, attenuate_cert = False).to(device_id)
@@ -189,11 +189,11 @@ def train(args):
     global_step = 0
     batch_size = args.gpu_batch_size
     step_size = gpus*batch_size
-    roma.STEP_SIZE = step_size
+    romatch.STEP_SIZE = step_size
     
     N = (32 * 250000)  # 250k steps of batch size 32
     # checkpoint every
-    k = 25000 // roma.STEP_SIZE
+    k = 25000 // romatch.STEP_SIZE
 
     # Data
     mega = MegadepthBuilder(data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True)
@@ -233,20 +233,20 @@ def train(args):
         alpha = 0.5,
         c = 1e-4,)
     parameters = [
-        {"params": model.encoder.parameters(), "lr": roma.STEP_SIZE * 5e-6 / 8},
-        {"params": model.decoder.parameters(), "lr": roma.STEP_SIZE * 1e-4 / 8},
+        {"params": model.encoder.parameters(), "lr": romatch.STEP_SIZE * 5e-6 / 8},
+        {"params": model.decoder.parameters(), "lr": romatch.STEP_SIZE * 1e-4 / 8},
     ]
     optimizer = torch.optim.AdamW(parameters, weight_decay=0.01)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer, milestones=[(9*N/roma.STEP_SIZE)//10])
+        optimizer, milestones=[(9*N/romatch.STEP_SIZE)//10])
     megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000, h=h,w=w)
     checkpointer = CheckPoint(checkpoint_dir, experiment_name)
     model, optimizer, lr_scheduler, global_step = checkpointer.load(model, optimizer, lr_scheduler, global_step)
-    roma.GLOBAL_STEP = global_step
+    romatch.GLOBAL_STEP = global_step
     ddp_model = DDP(model, device_ids=[device_id], find_unused_parameters = False, gradient_as_bucket_view=True)
     grad_scaler = torch.cuda.amp.GradScaler(growth_interval=1_000_000)
     grad_clip_norm = 0.01
-    for n in range(roma.GLOBAL_STEP, N, k * roma.STEP_SIZE):
+    for n in range(romatch.GLOBAL_STEP, N, k * romatch.STEP_SIZE):
         mega_sampler = torch.utils.data.WeightedRandomSampler(
             mega_ws, num_samples = batch_size * k, replacement=False
         )
@@ -269,15 +269,15 @@ def train(args):
                 num_workers=gpus * 8,
             )
         )
-        for n_k in tqdm(range(n, n + 2 * k, 2),disable = roma.RANK > 0):
+        for n_k in tqdm(range(n, n + 2 * k, 2),disable = romatch.RANK > 0):
             train_k_steps(
                 n_k, 1, mega_dataloader, ddp_model, depth_loss_mega, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm, progress_bar=False
             )
             train_k_steps(
                 n_k + 1, 1, scannet_dataloader, ddp_model, depth_loss_scannet, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm, progress_bar=False
             )
-        checkpointer.save(model, optimizer, lr_scheduler, roma.GLOBAL_STEP)
-        wandb.log(megadense_benchmark.benchmark(model), step = roma.GLOBAL_STEP)
+        checkpointer.save(model, optimizer, lr_scheduler, romatch.GLOBAL_STEP)
+        wandb.log(megadense_benchmark.benchmark(model), step = romatch.GLOBAL_STEP)
 
 def test_scannet(model, name, resolution, sample_mode):
     scannet_benchmark = ScanNetBenchmark("data/scannet")
@@ -291,7 +291,7 @@ if __name__ == "__main__":
     os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
     os.environ["OMP_NUM_THREADS"] = "16"
     
-    import roma
+    import romatch
     parser = ArgumentParser()
     parser.add_argument("--test", action='store_true')
     parser.add_argument("--debug_mode", action='store_true')
@@ -301,7 +301,7 @@ if __name__ == "__main__":
     parser.add_argument("--wandb_entity", required = False)
 
     args, _ = parser.parse_known_args()
-    roma.DEBUG_MODE = args.debug_mode
+    romatch.DEBUG_MODE = args.debug_mode
     if not args.test:
         train(args)
     experiment_name = os.path.splitext(os.path.basename(__file__))[0]

experiments/train_roma_outdoor.py

@@ -9,16 +9,16 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 import json
 import wandb
 
-from roma.benchmarks import MegadepthDenseBenchmark
-from roma.datasets.megadepth import MegadepthBuilder
-from roma.losses.robust_loss import RobustLosses
-from roma.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
+from romatch.benchmarks import MegadepthDenseBenchmark
+from romatch.datasets.megadepth import MegadepthBuilder
+from romatch.losses.robust_loss import RobustLosses
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
 
-from roma.train.train import train_k_steps
-from roma.models.matcher import *
-from roma.models.transformer import Block, TransformerDecoder, MemEffAttention
-from roma.models.encoders import *
-from roma.checkpointing import CheckPoint
+from romatch.train.train import train_k_steps
+from romatch.models.matcher import *
+from romatch.models.transformer import Block, TransformerDecoder, MemEffAttention
+from romatch.models.encoders import *
+from romatch.checkpointing import CheckPoint
 
 resolutions = {"low":(448, 448), "medium":(14*8*5, 14*8*5), "high":(14*8*6, 14*8*6)}
 
@@ -174,14 +174,14 @@ def train(args):
     rank = dist.get_rank()
     print(f"Start running DDP on rank {rank}")
     device_id = rank % torch.cuda.device_count()
-    roma.LOCAL_RANK = device_id
+    romatch.LOCAL_RANK = device_id
     torch.cuda.set_device(device_id)
     
     resolution = args.train_resolution
     wandb_log = not args.dont_log_wandb
     experiment_name = os.path.splitext(os.path.basename(__file__))[0]
     wandb_mode = "online" if wandb_log and rank == 0 else "disabled"
-    wandb.init(project="roma", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
+    wandb.init(project="romatch", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
     checkpoint_dir = "workspace/checkpoints/"
     h,w = resolutions[resolution]
     model = get_model(pretrained_backbone=True, resolution=resolution, attenuate_cert = False).to(device_id)
@@ -189,11 +189,11 @@ def train(args):
     global_step = 0
     batch_size = args.gpu_batch_size
     step_size = gpus*batch_size
-    roma.STEP_SIZE = step_size
+    romatch.STEP_SIZE = step_size
     
     N = (32 * 250000)  # 250k steps of batch size 32
     # checkpoint every
-    k = 25000 // roma.STEP_SIZE
+    k = 25000 // romatch.STEP_SIZE
 
     # Data
     mega = MegadepthBuilder(data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True)
@@ -219,20 +219,20 @@ def train(args):
         alpha = 0.5,
         c = 1e-4,)
     parameters = [
-        {"params": model.encoder.parameters(), "lr": roma.STEP_SIZE * 5e-6 / 8},
-        {"params": model.decoder.parameters(), "lr": roma.STEP_SIZE * 1e-4 / 8},
+        {"params": model.encoder.parameters(), "lr": romatch.STEP_SIZE * 5e-6 / 8},
+        {"params": model.decoder.parameters(), "lr": romatch.STEP_SIZE * 1e-4 / 8},
     ]
     optimizer = torch.optim.AdamW(parameters, weight_decay=0.01)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer, milestones=[(9*N/roma.STEP_SIZE)//10])
+        optimizer, milestones=[(9*N/romatch.STEP_SIZE)//10])
     megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000, h=h,w=w)
     checkpointer = CheckPoint(checkpoint_dir, experiment_name)
     model, optimizer, lr_scheduler, global_step = checkpointer.load(model, optimizer, lr_scheduler, global_step)
-    roma.GLOBAL_STEP = global_step
+    romatch.GLOBAL_STEP = global_step
     ddp_model = DDP(model, device_ids=[device_id], find_unused_parameters = False, gradient_as_bucket_view=True)
     grad_scaler = torch.cuda.amp.GradScaler(growth_interval=1_000_000)
     grad_clip_norm = 0.01
-    for n in range(roma.GLOBAL_STEP, N, k * roma.STEP_SIZE):
+    for n in range(romatch.GLOBAL_STEP, N, k * romatch.STEP_SIZE):
         mega_sampler = torch.utils.data.WeightedRandomSampler(
             mega_ws, num_samples = batch_size * k, replacement=False
         )
@@ -247,8 +247,8 @@ def train(args):
         train_k_steps(
             n, k, mega_dataloader, ddp_model, depth_loss, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm,
         )
-        checkpointer.save(model, optimizer, lr_scheduler, roma.GLOBAL_STEP)
-        wandb.log(megadense_benchmark.benchmark(model), step = roma.GLOBAL_STEP)
+        checkpointer.save(model, optimizer, lr_scheduler, romatch.GLOBAL_STEP)
+        wandb.log(megadense_benchmark.benchmark(model), step = romatch.GLOBAL_STEP)
 
 def test_mega_8_scenes(model, name):
     mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
@@ -292,7 +292,7 @@ if __name__ == "__main__":
     os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
     os.environ["OMP_NUM_THREADS"] = "16"
     torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
-    import roma
+    import romatch
     parser = ArgumentParser()
     parser.add_argument("--only_test", action='store_true')
     parser.add_argument("--debug_mode", action='store_true')
@@ -302,6 +302,6 @@ if __name__ == "__main__":
     parser.add_argument("--wandb_entity", required = False)
 
     args, _ = parser.parse_known_args()
-    roma.DEBUG_MODE = args.debug_mode
+    romatch.DEBUG_MODE = args.debug_mode
     if not args.only_test:
         train(args)

experiments/train_tiny_roma_v1_outdoor.py

@@ -17,13 +17,13 @@ import wandb
 from PIL import Image
 from torchvision.transforms import ToTensor
 
-from roma.benchmarks import MegadepthDenseBenchmark, ScanNetBenchmark
-from roma.benchmarks import Mega1500PoseLibBenchmark, ScanNetPoselibBenchmark
-from roma.datasets.megadepth import MegadepthBuilder
-from roma.losses.robust_loss_tiny_roma import RobustLosses
-from roma.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
-from roma.train.train import train_k_steps
-from roma.checkpointing import CheckPoint
+from romatch.benchmarks import MegadepthDenseBenchmark, ScanNetBenchmark
+from romatch.benchmarks import Mega1500PoseLibBenchmark, ScanNetPoselibBenchmark
+from romatch.datasets.megadepth import MegadepthBuilder
+from romatch.losses.robust_loss_tiny_roma import RobustLosses
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
+from romatch.train.train import train_k_steps
+from romatch.checkpointing import CheckPoint
 
 resolutions = {"low":(448, 448), "medium":(14*8*5, 14*8*5), "high":(14*8*6, 14*8*6), "xfeat": (600,800), "big": (768, 1024)}
 
@@ -194,7 +194,7 @@ class XFeatModel(nn.Module):
             white_im = torch.ones((H, W), device = device)
         vis_im = certainty * warp_im + (1 - certainty) * white_im
         if save_path is not None:
-            from roma.utils import tensor_to_pil
+            from romatch.utils import tensor_to_pil
             tensor_to_pil(vis_im, unnormalize=unnormalize).save(save_path)
         return vis_im
      
@@ -330,14 +330,14 @@ def train(args):
     rank = 0
     gpus = 1
     device_id = rank % torch.cuda.device_count()
-    roma.LOCAL_RANK = 0
+    romatch.LOCAL_RANK = 0
     torch.cuda.set_device(device_id)
         
     resolution = "big"
     wandb_log = not args.dont_log_wandb
     experiment_name = Path(__file__).stem
     wandb_mode = "online" if wandb_log and rank == 0 else "disabled"
-    wandb.init(project="roma", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
+    wandb.init(project="romatch", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
     checkpoint_dir = "workspace/checkpoints/"
     h,w = resolutions[resolution]
     model = XFeatModel(freeze_xfeat = False).to(device_id)
@@ -345,11 +345,11 @@ def train(args):
     global_step = 0
     batch_size = args.gpu_batch_size
     step_size = gpus*batch_size
-    roma.STEP_SIZE = step_size
+    romatch.STEP_SIZE = step_size
     
     N = 2_000_000  # 2M pairs
     # checkpoint every
-    k = 25000 // roma.STEP_SIZE
+    k = 25000 // romatch.STEP_SIZE
 
     # Data
     mega = MegadepthBuilder(data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True)
@@ -377,21 +377,21 @@ def train(args):
         epe_mask_prob_th = 0.001,
         )
     parameters = [
-        {"params": model.parameters(), "lr": roma.STEP_SIZE * 1e-4 / 8},
+        {"params": model.parameters(), "lr": romatch.STEP_SIZE * 1e-4 / 8},
     ]
     optimizer = torch.optim.AdamW(parameters, weight_decay=0.01)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer, milestones=[(9*N/roma.STEP_SIZE)//10])
+        optimizer, milestones=[(9*N/romatch.STEP_SIZE)//10])
     #megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000, h=h,w=w)
     mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth", num_ransac_iter = 1, test_every = 30)
 
     checkpointer = CheckPoint(checkpoint_dir, experiment_name)
     model, optimizer, lr_scheduler, global_step = checkpointer.load(model, optimizer, lr_scheduler, global_step)
-    roma.GLOBAL_STEP = global_step
+    romatch.GLOBAL_STEP = global_step
     grad_scaler = torch.cuda.amp.GradScaler(growth_interval=1_000_000)
     grad_clip_norm = 0.01
     #megadense_benchmark.benchmark(model)
-    for n in range(roma.GLOBAL_STEP, N, k * roma.STEP_SIZE):
+    for n in range(romatch.GLOBAL_STEP, N, k * romatch.STEP_SIZE):
         mega_sampler = torch.utils.data.WeightedRandomSampler(
             mega_ws, num_samples = batch_size * k, replacement=False
         )
@@ -406,8 +406,8 @@ def train(args):
         train_k_steps(
             n, k, mega_dataloader, model, depth_loss, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm,
         )
-        checkpointer.save(model, optimizer, lr_scheduler, roma.GLOBAL_STEP)
-        wandb.log(mega1500_benchmark.benchmark(model, model_name=experiment_name), step = roma.GLOBAL_STEP)
+        checkpointer.save(model, optimizer, lr_scheduler, romatch.GLOBAL_STEP)
+        wandb.log(mega1500_benchmark.benchmark(model, model_name=experiment_name), step = romatch.GLOBAL_STEP)
 
 def test_mega_8_scenes(model, name):
     mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
@@ -476,7 +476,7 @@ if __name__ == "__main__":
     os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
     os.environ["OMP_NUM_THREADS"] = "16"
     torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
-    import roma
+    import romatch
     parser = ArgumentParser()
     parser.add_argument("--only_test", action='store_true')
     parser.add_argument("--debug_mode", action='store_true')
@@ -486,7 +486,7 @@ if __name__ == "__main__":
     parser.add_argument("--wandb_entity", required = False)
 
     args, _ = parser.parse_known_args()
-    roma.DEBUG_MODE = args.debug_mode
+    romatch.DEBUG_MODE = args.debug_mode
     if not args.only_test:
         train(args)
 

roma/benchmarks/hpatches_sequences_homog_benchmark.py → romatch/benchmarks/hpatches_sequences_homog_benchmark.py

@@ -4,7 +4,7 @@ import numpy as np
 import os
 
 from tqdm import tqdm
-from roma.utils import pose_auc
+from romatch.utils import pose_auc
 import cv2
 
 

roma/benchmarks/megadepth_dense_benchmark.py → romatch/benchmarks/megadepth_dense_benchmark.py

@@ -1,10 +1,10 @@
 import torch
 import numpy as np
 import tqdm
-from roma.datasets import MegadepthBuilder
-from roma.utils import warp_kpts
+from romatch.datasets import MegadepthBuilder
+from romatch.utils import warp_kpts
 from torch.utils.data import ConcatDataset
-import roma
+import romatch
 
 class MegadepthDenseBenchmark:
     def __init__(self, data_root="data/megadepth", h = 384, w = 512, num_samples = 2000) -> None:
@@ -55,7 +55,7 @@ class MegadepthDenseBenchmark:
             dataloader = torch.utils.data.DataLoader(
                 self.dataset, batch_size=B, num_workers=batch_size, sampler=sampler
             )
-            for idx, data in tqdm.tqdm(enumerate(dataloader), disable = roma.RANK > 0):
+            for idx, data in tqdm.tqdm(enumerate(dataloader), disable = romatch.RANK > 0):
                 im_A, im_B, depth1, depth2, T_1to2, K1, K2 = (
                     data["im_A"].cuda(),
                     data["im_B"].cuda(),
@@ -69,8 +69,8 @@ class MegadepthDenseBenchmark:
                 gd, pck_1, pck_3, pck_5, prob = self.geometric_dist(
                     depth1, depth2, T_1to2, K1, K2, matches
                 )
-                if roma.DEBUG_MODE:
-                    from roma.utils.utils import tensor_to_pil
+                if romatch.DEBUG_MODE:
+                    from romatch.utils.utils import tensor_to_pil
                     import torch.nn.functional as F
                     path = "vis"
                     H, W = model.get_output_resolution()

roma/benchmarks/megadepth_pose_estimation_benchmark.py → romatch/benchmarks/megadepth_pose_estimation_benchmark.py

@@ -1,10 +1,10 @@
 import numpy as np
 import torch
-from roma.utils import *
+from romatch.utils import *
 from PIL import Image
 from tqdm import tqdm
 import torch.nn.functional as F
-import roma
+import romatch
 import kornia.geometry.epipolar as kepi
 
 class MegaDepthPoseEstimationBenchmark:

roma/benchmarks/megadepth_pose_estimation_benchmark_poselib.py → romatch/benchmarks/megadepth_pose_estimation_benchmark_poselib.py

@@ -1,10 +1,10 @@
 import numpy as np
 import torch
-from roma.utils import *
+from romatch.utils import *
 from PIL import Image
 from tqdm import tqdm
 import torch.nn.functional as F
-import roma
+import romatch
 import kornia.geometry.epipolar as kepi
 
 # wrap cause pyposelib is still in dev

roma/benchmarks/scannet_benchmark.py → romatch/benchmarks/scannet_benchmark.py

@@ -1,7 +1,7 @@
 import os.path as osp
 import numpy as np
 import torch
-from roma.utils import *
+from romatch.utils import *
 from PIL import Image
 from tqdm import tqdm
 

roma/checkpointing/checkpoint.py → romatch/checkpointing/checkpoint.py

@@ -5,7 +5,7 @@ from torch.nn.parallel.distributed import DistributedDataParallel
 from loguru import logger
 import gc
 
-import roma
+import romatch
 
 class CheckPoint:
     def __init__(self, dir=None, name="tmp"):
@@ -20,7 +20,7 @@ class CheckPoint:
         lr_scheduler,
         n,
         ):
-        if roma.RANK == 0:
+        if romatch.RANK == 0:
             assert model is not None
             if isinstance(model, (DataParallel, DistributedDataParallel)):
                 model = model.module
@@ -40,7 +40,7 @@ class CheckPoint:
         lr_scheduler,
         n,
         ):
-        if os.path.exists(self.dir + self.name + f"_latest.pth") and roma.RANK == 0:
+        if os.path.exists(self.dir + self.name + f"_latest.pth") and romatch.RANK == 0:
             states = torch.load(self.dir + self.name + f"_latest.pth")
             if "model" in states:
                 model.load_state_dict(states["model"])

roma/datasets/megadepth.py → romatch/datasets/megadepth.py

@@ -5,9 +5,9 @@ import numpy as np
 import torch
 import torchvision.transforms.functional as tvf
 import kornia.augmentation as K
-from roma.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
-import roma
-from roma.utils import *
+from romatch.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
+import romatch
+from romatch.utils import *
 import math
 
 class MegadepthScene:
@@ -53,7 +53,7 @@ class MegadepthScene:
             area = ht * wt
             s = int(16 * (math.sqrt(area)//16))
             sizes = ((ht,wt), (s,s), (wt,ht))
-            choice = roma.RANK % 3
+            choice = romatch.RANK % 3
             ht, wt = sizes[choice] 
         # counts, bins = np.histogram(self.overlaps,20)
         # print(counts)
@@ -157,7 +157,7 @@ class MegadepthScene:
             if np.random.rand() > 0.5:
                 im_B, depth_B, K2 = self.single_horizontal_flip(im_B, depth_B, K2)
         
-        if roma.DEBUG_MODE:
+        if romatch.DEBUG_MODE:
             tensor_to_pil(im_A[0], unnormalize=True).save(
                             f"vis/im_A.jpg")
             tensor_to_pil(im_B[0], unnormalize=True).save(

roma/datasets/scannet.py → romatch/datasets/scannet.py

@@ -14,9 +14,9 @@ import torchvision.transforms.functional as tvf
 import kornia.augmentation as K
 import os.path as osp
 import matplotlib.pyplot as plt
-import roma
-from roma.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
-from roma.utils.transforms import GeometricSequential
+import romatch
+from romatch.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
+from romatch.utils.transforms import GeometricSequential
 from tqdm import tqdm
 
 class ScanNetScene:
@@ -147,7 +147,7 @@ class ScanNetBuilder:
         # Note: split doesn't matter here as we always use same scannet_train scenes
         scene_names = self.all_scenes
         scenes = []
-        for scene_name in tqdm(scene_names, disable = roma.RANK > 0):
+        for scene_name in tqdm(scene_names, disable = romatch.RANK > 0):
             scene_info = np.load(os.path.join(self.scene_info_root,scene_name), allow_pickle=True)
             scenes.append(ScanNetScene(self.data_root, scene_info, min_overlap=min_overlap, **kwargs))
         return scenes

roma/losses/robust_loss.py → romatch/losses/robust_loss.py

@@ -2,9 +2,9 @@ from einops.einops import rearrange
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from roma.utils.utils import get_gt_warp
+from romatch.utils.utils import get_gt_warp
 import wandb
-import roma
+import romatch
 import math
 
 class RobustLosses(nn.Module):
@@ -57,7 +57,7 @@ class RobustLosses(nn.Module):
             f"gm_certainty_loss_{scale}": certainty_loss.mean(),
             f"gm_cls_loss_{scale}": cls_loss.mean(),
         }
-        wandb.log(losses, step = roma.GLOBAL_STEP)
+        wandb.log(losses, step = romatch.GLOBAL_STEP)
         return losses
 
     def delta_cls_loss(self, x2, prob, flow_pre_delta, delta_cls, certainty, scale, offset_scale):
@@ -76,14 +76,14 @@ class RobustLosses(nn.Module):
             f"delta_certainty_loss_{scale}": certainty_loss.mean(),
             f"delta_cls_loss_{scale}": cls_loss.mean(),
         }
-        wandb.log(losses, step = roma.GLOBAL_STEP)
+        wandb.log(losses, step = romatch.GLOBAL_STEP)
         return losses
 
     def regression_loss(self, x2, prob, flow, certainty, scale, eps=1e-8, mode = "delta"):
         epe = (flow.permute(0,2,3,1) - x2).norm(dim=-1)
         if scale == 1:
             pck_05 = (epe[prob > 0.99] < 0.5 * (2/512)).float().mean()
-            wandb.log({"train_pck_05": pck_05}, step = roma.GLOBAL_STEP)
+            wandb.log({"train_pck_05": pck_05}, step = romatch.GLOBAL_STEP)
 
         ce_loss = F.binary_cross_entropy_with_logits(certainty[:, 0], prob)
         a = self.alpha[scale] if isinstance(self.alpha, dict) else self.alpha
@@ -96,7 +96,7 @@ class RobustLosses(nn.Module):
             f"{mode}_certainty_loss_{scale}": ce_loss.mean(),
             f"{mode}_regression_loss_{scale}": reg_loss.mean(),
         }
-        wandb.log(losses, step = roma.GLOBAL_STEP)
+        wandb.log(losses, step = romatch.GLOBAL_STEP)
         return losses
 
     def forward(self, corresps, batch):

roma/losses/robust_loss_tiny_roma.py → romatch/losses/robust_loss_tiny_roma.py

@@ -2,12 +2,12 @@ from einops.einops import rearrange
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from roma.utils.utils import get_gt_warp
+from romatch.utils.utils import get_gt_warp
 import wandb
-import roma
+import romatch
 import math
 
-# This is slightly different than regular roma due to significantly worse corresps
+# This is slightly different than regular romatch due to significantly worse corresps
 # The confidence loss is quite tricky here //Johan
 
 class RobustLosses(nn.Module):
@@ -57,7 +57,7 @@ class RobustLosses(nn.Module):
         losses = {
             f"gm_corr_volume_loss_{scale}": corr_volume_loss.mean(),
         }
-        wandb.log(losses, step = roma.GLOBAL_STEP)
+        wandb.log(losses, step = romatch.GLOBAL_STEP)
         return losses
 
     
@@ -68,7 +68,7 @@ class RobustLosses(nn.Module):
             prob = prob * (epe < (2 / 512) * (self.local_dist[scale] * scale)).float()
         if scale == 1:
             pck_05 = (epe[prob > 0.99] < 0.5 * (2/512)).float().mean()
-            wandb.log({"train_pck_05": pck_05}, step = roma.GLOBAL_STEP)
+            wandb.log({"train_pck_05": pck_05}, step = romatch.GLOBAL_STEP)
         if self.epe_mask_prob_th is not None:
             # if too far away from gt, certainty should be 0
             gt_cert = prob * (epe < scale * self.epe_mask_prob_th)
@@ -88,7 +88,7 @@ class RobustLosses(nn.Module):
             f"{mode}_certainty_loss_{scale}": ce_loss.mean(),
             f"{mode}_regression_loss_{scale}": reg_loss.mean(),
         }
-        wandb.log(losses, step = roma.GLOBAL_STEP)
+        wandb.log(losses, step = romatch.GLOBAL_STEP)
         return losses
 
     def forward(self, corresps, batch):

roma/models/matcher.py → romatch/models/matcher.py

@@ -9,11 +9,11 @@ import warnings
 from warnings import warn
 from PIL import Image
 
-import roma
-from roma.utils import get_tuple_transform_ops
-from roma.utils.local_correlation import local_correlation
-from roma.utils.utils import cls_to_flow_refine
-from roma.utils.kde import kde
+import romatch
+from romatch.utils import get_tuple_transform_ops
+from romatch.utils.local_correlation import local_correlation
+from romatch.utils.utils import cls_to_flow_refine
+from romatch.utils.kde import kde
 from typing import Union
 
 class ConvRefiner(nn.Module):
@@ -767,6 +767,6 @@ class RegressionMatcher(nn.Module):
             white_im = torch.ones((H, W), device = device)
         vis_im = certainty * warp_im + (1 - certainty) * white_im
         if save_path is not None:
-            from roma.utils import tensor_to_pil
+            from romatch.utils import tensor_to_pil
             tensor_to_pil(vis_im, unnormalize=unnormalize).save(save_path)
         return vis_im

roma/models/model_zoo/__init__.py → romatch/models/model_zoo/__init__.py

@@ -3,12 +3,12 @@ import torch
 from .roma_models import roma_model, tiny_roma_v1_model
 
 weight_urls = {
-    "roma": {
-        "outdoor": "https://github.com/Parskatt/storage/releases/download/roma/roma_outdoor.pth",
-        "indoor": "https://github.com/Parskatt/storage/releases/download/roma/roma_indoor.pth",
+    "romatch": {
+        "outdoor": "https://github.com/Parskatt/storage/releases/download/romatch/roma_outdoor.pth",
+        "indoor": "https://github.com/Parskatt/storage/releases/download/romatch/roma_indoor.pth",
     },
     "tiny_roma_v1": {
-        "outdoor": "https://github.com/Parskatt/storage/releases/download/roma/tiny_roma_v1_outdoor.pth",
+        "outdoor": "https://github.com/Parskatt/storage/releases/download/romatch/tiny_roma_v1_outdoor.pth",
     },
     "dinov2": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth", #hopefully this doesnt change :D
 }
@@ -37,7 +37,7 @@ def roma_outdoor(device, weights=None, dinov2_weights=None, coarse_res: Union[in
     assert coarse_res[1] % 14 == 0, "Needs to be multiple of 14 for backbone"
     
     if weights is None:
-        weights = torch.hub.load_state_dict_from_url(weight_urls["roma"]["outdoor"],
+        weights = torch.hub.load_state_dict_from_url(weight_urls["romatch"]["outdoor"],
                                                      map_location=device)
     if dinov2_weights is None:
         dinov2_weights = torch.hub.load_state_dict_from_url(weight_urls["dinov2"],
@@ -58,7 +58,7 @@ def roma_indoor(device, weights=None, dinov2_weights=None, coarse_res: Union[int
     assert coarse_res[1] % 14 == 0, "Needs to be multiple of 14 for backbone"
     
     if weights is None:
-        weights = torch.hub.load_state_dict_from_url(weight_urls["roma"]["indoor"],
+        weights = torch.hub.load_state_dict_from_url(weight_urls["romatch"]["indoor"],
                                                      map_location=device)
     if dinov2_weights is None:
         dinov2_weights = torch.hub.load_state_dict_from_url(weight_urls["dinov2"],

roma/models/model_zoo/roma_models.py → romatch/models/model_zoo/roma_models.py

@@ -1,10 +1,10 @@
 import warnings
 import torch.nn as nn
 import torch
-from roma.models.matcher import *
-from roma.models.transformer import Block, TransformerDecoder, MemEffAttention
-from roma.models.encoders import *
-from roma.models.tiny import TinyRoMa
+from romatch.models.matcher import *
+from romatch.models.transformer import Block, TransformerDecoder, MemEffAttention
+from romatch.models.encoders import *
+from romatch.models.tiny import TinyRoMa
 
 def tiny_roma_v1_model(weights = None, freeze_xfeat=False, exact_softmax=False, xfeat = None):
     model = TinyRoMa(
@@ -16,7 +16,7 @@ def tiny_roma_v1_model(weights = None, freeze_xfeat=False, exact_softmax=False,
     return model
 
 def roma_model(resolution, upsample_preds, device = None, weights=None, dinov2_weights=None, amp_dtype: torch.dtype=torch.float16, **kwargs):
-    # roma weights and dinov2 weights are loaded seperately, as dinov2 weights are not parameters
+    # romatch weights and dinov2 weights are loaded seperately, as dinov2 weights are not parameters
     #torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul TODO: these probably ruin stuff, should be careful
     #torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
     warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')

roma/models/tiny.py → romatch/models/tiny.py

@@ -10,7 +10,7 @@ import numpy as np
 from torch import nn
 from PIL import Image
 from torchvision.transforms import ToTensor
-from roma.utils.kde import kde
+from romatch.utils.kde import kde
 
 class BasicLayer(nn.Module):
     """
@@ -171,7 +171,7 @@ class TinyRoMa(nn.Module):
             white_im = torch.ones((H, W), device = device)
         vis_im = certainty * warp_im + (1 - certainty) * white_im
         if save_path is not None:
-            from roma.utils import tensor_to_pil
+            from romatch.utils import tensor_to_pil
             tensor_to_pil(vis_im, unnormalize=unnormalize).save(save_path)
         return vis_im
      

roma/models/transformer/__init__.py → romatch/models/transformer/__init__.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from roma.utils.utils import get_grid
+from romatch.utils.utils import get_grid
 from .layers.block import Block
 from .layers.attention import MemEffAttention
 from .dinov2 import vit_large

roma/train/train.py → romatch/train/train.py

@@ -1,6 +1,6 @@
 from tqdm import tqdm
-from roma.utils.utils import to_cuda
-import roma
+from romatch.utils.utils import to_cuda
+import romatch
 import torch
 import wandb
 
@@ -17,8 +17,8 @@ def log_param_statistics(named_parameters, norm_type = 2):
     total_grad_norm = torch.norm(grad_norms, norm_type)
     if torch.any(nans_or_infs):
         print(f"These params have nan or inf grads: {nan_inf_names}")
-    wandb.log({"grad_norm": total_grad_norm.item()}, step = roma.GLOBAL_STEP)
-    wandb.log({"param_norm": param_norm.item()}, step = roma.GLOBAL_STEP)
+    wandb.log({"grad_norm": total_grad_norm.item()}, step = romatch.GLOBAL_STEP)
+    wandb.log({"param_norm": param_norm.item()}, step = romatch.GLOBAL_STEP)
 
 def train_step(train_batch, model, objective, optimizer, grad_scaler, grad_clip_norm = 1.,**kwargs):
     optimizer.zero_grad()
@@ -30,17 +30,17 @@ def train_step(train_batch, model, objective, optimizer, grad_scaler, grad_clip_
     torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip_norm) # what should max norm be?
     grad_scaler.step(optimizer)
     grad_scaler.update()
-    wandb.log({"grad_scale": grad_scaler._scale.item()}, step = roma.GLOBAL_STEP)
+    wandb.log({"grad_scale": grad_scaler._scale.item()}, step = romatch.GLOBAL_STEP)
     if grad_scaler._scale < 1.:
         grad_scaler._scale = torch.tensor(1.).to(grad_scaler._scale)
-    roma.GLOBAL_STEP = roma.GLOBAL_STEP + roma.STEP_SIZE # increment global step
+    romatch.GLOBAL_STEP = romatch.GLOBAL_STEP + romatch.STEP_SIZE # increment global step
     return {"train_out": out, "train_loss": l.item()}
 
 
 def train_k_steps(
     n_0, k, dataloader, model, objective, optimizer, lr_scheduler, grad_scaler, progress_bar=True, grad_clip_norm = 1., warmup = None, ema_model = None, pbar_n_seconds = 1,
 ):
-    for n in tqdm(range(n_0, n_0 + k), disable=(not progress_bar) or roma.RANK > 0, mininterval=pbar_n_seconds):
+    for n in tqdm(range(n_0, n_0 + k), disable=(not progress_bar) or romatch.RANK > 0, mininterval=pbar_n_seconds):
         batch = next(dataloader)
         model.train(True)
         batch = to_cuda(batch)

setup.py

@@ -1,8 +1,8 @@
 from setuptools import setup, find_packages
 
 setup(
-    name="roma",
-    packages=find_packages(include=("roma*",)),
+    name="romatch",
+    packages=find_packages(include=("romatch*",)),
     version="0.0.1",
     author="Johan Edstedt",
     install_requires=open("requirements.txt", "r").read().split("\n"),