RomaFeaturesMatcher: support multiple features folders

meshroom/roma/MatchMasking.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from meshroom.core import desc
 from meshroom.core.utils import DESCRIBER_TYPES
+from pyalicevision import parallelization as avpar
 
 class MatchMasking(desc.CommandLineNode):
 
@@ -12,7 +13,7 @@ class MatchMasking(desc.CommandLineNode):
     Mask the certainty images generated by RomaMatcher given the associated 
     masks of the reference and matched images.
     """
-    size = desc.DynamicNodeSize('inputSfMData')
+    size = avpar.DynamicViewsSize('inputSfMData')
 
     parallelization = desc.Parallelization(blockSize=40)
     commandLineRange = "--rangeIteration {rangeIteration} --rangeBlocksCount {rangeBlocksCount}"

meshroom/roma/RomaFeaturesMatcher.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 from meshroom.core import desc
+from pyalicevision import parallelization as avpar
 
 class RomaFeaturesMatcher(desc.CommandLineNode):
 
@@ -14,7 +15,7 @@ class RomaFeaturesMatcher(desc.CommandLineNode):
     A region will only be matched to features in the other image given the warp image coordinates.
     """
 
-    size = desc.DynamicNodeSize('inputSfMData')
+    size = avpar.DynamicViewsSize('inputSfMData')
     cpu = desc.Level.INTENSIVE
 
     parallelization = desc.Parallelization(blockSize=40)
@@ -32,11 +33,16 @@ class RomaFeaturesMatcher(desc.CommandLineNode):
             description="Input SfMData file.",
             value="",
         ),
-        desc.File(
-            name="featuresFolder",
-            label="Features folder",
-            description="Input features",
-            value=""
+        desc.ListAttribute(
+            elementDesc=desc.File(
+                name="featuresFolder",
+                label="Features Folder",
+                description="Folder containing input features.",
+                value="",
+            ),
+            name="featuresFolders",
+            label="Features Folders",
+            description="Input features folders.",
         ),
         desc.File(
             name="imagePairsList",

meshroom/roma/RomaMatcher.py

@@ -2,16 +2,16 @@
 
 import os 
 from pathlib import Path
-
 from meshroom.core import desc
+from pyalicevision import parallelization as avpar
 
 class RomaMatcher(desc.CommandLineNode):
 
     category = "ROMA"
     documentation = """
     Compute ROMA warp and certainty images on a list of images pairs.
     """
-    size = desc.DynamicNodeSize('inputSfMData')
+    size = avpar.DynamicViewsSize('inputSfMData')
     gpu = desc.Level.INTENSIVE
 
     parallelization = desc.Parallelization(blockSize=40)

meshroom/roma/RomaReducer.py

@@ -10,7 +10,6 @@ class RomaReducer(desc.CommandLineNode):
     category = "ROMA"
     documentation = """ RomaSampler is parallelized and distributed. It generated multiple
     output files which must merged together. """
-    size = desc.DynamicNodeSize('inputSfMData')
 
     exePath = (Path(__file__).absolute().parent.parent.parent / "python" / "reducer.py").as_posix()
     commandLine="python "+exePath+" {allParams}"
@@ -42,7 +41,7 @@ class RomaReducer(desc.CommandLineNode):
             value=["sift"],
             exclusive=False,
             joinChar=",",
-            group="ingored"
+            commandLineGroup="ignored"
         ),
     ]
 

meshroom/roma/RomaSampler.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from meshroom.core import desc
 from meshroom.core.utils import DESCRIBER_TYPES
+from pyalicevision import parallelization as avpar
+
 
 class RomaSampler(desc.CommandLineNode):
 
@@ -12,7 +14,7 @@ class RomaSampler(desc.CommandLineNode):
     Sample the dense ROMA matches to generate features/matches used by SFM.
     This is an intermediate node which has to be followed by RomaReducer.
     """
-    size = desc.DynamicNodeSize('inputSfMData')
+    size = avpar.DynamicViewsSize('inputSfMData')
 
     parallelization = desc.Parallelization(blockSize=40)
     commandLineRange = "--rangeIteration {rangeIteration} --rangeBlocksCount {rangeBlocksCount}"
@@ -79,7 +81,7 @@ class RomaSampler(desc.CommandLineNode):
             value=["sift"],
             exclusive=False,
             joinChar=",",
-            group="ingored"
+            commandLineGroup="ignored"
         ),
     ]
 

meshroom/roma/StarListing.py

@@ -13,7 +13,6 @@ class StarListing(desc.Node):
     Assume the current keyframe is keyframe #N, all the frames between #N and #N+radiusKeyFrames will be matched.
     Same for all the frames between #N and #N-radiusKeyFrames.
     """
-    size = desc.DynamicNodeSize("inputSfMData")
 
     inputs = [
         desc.File(

python/featuresMatcher.py

@@ -23,7 +23,7 @@ def regionToNumpy(region):
 
     return array
 
-def compute_featuresMatcher(inputSfMData, imagePairsList, warpFolder, featuresFolder, matchesFolder, masksFolder, masksExtension, minCertainty, rangeIteration, rangeBlocksCount):
+def compute_featuresMatcher(inputSfMData, imagePairsList, warpFolder, featuresFolders, matchesFolder, masksFolder, masksExtension, minCertainty, rangeIteration, rangeBlocksCount):
 
      # Parse sfm
     iinfos = get_imageinfos_from_sfmdata(inputSfMData)
@@ -54,8 +54,18 @@ def compute_featuresMatcher(inputSfMData, imagePairsList, warpFolder, featuresFo
         #load features
         regionsRef = avfeat.SiftRegions()
         regionsOther = avfeat.SiftRegions()
-        regionsRef.Load(f"{featuresFolder}/{referenceId}.dspsift.feat", f"{featuresFolder}/{referenceId}.dspsift.desc")
-        regionsOther.Load(f"{featuresFolder}/{otherId}.dspsift.feat", f"{featuresFolder}/{otherId}.dspsift.desc")
+        for folder in featuresFolders:
+            ref_feat = f"{folder}/{referenceId}.dspsift.feat"
+            ref_desc = f"{folder}/{referenceId}.dspsift.desc"
+            if os.path.exists(ref_feat) and os.path.exists(ref_desc):
+                regionsRef.Load(ref_feat, ref_desc)
+                break
+        for folder in featuresFolders:
+            other_feat = f"{folder}/{otherId}.dspsift.feat"
+            other_desc = f"{folder}/{otherId}.dspsift.desc"
+            if os.path.exists(other_feat) and os.path.exists(other_desc):
+                regionsOther.Load(other_feat, other_desc)
+                break
 
         #load warp
         pair_string = str(referenceId) + "_" + str(otherId)
@@ -181,7 +191,7 @@ def compute_featuresMatcher(inputSfMData, imagePairsList, warpFolder, featuresFo
     parser.add_argument('--inputSfMData', type=str, help='')
     parser.add_argument('--imagePairsList', type=str, help='')
     parser.add_argument('--warpFolder', type=str, help='')
-    parser.add_argument('--featuresFolder', type=str, help='')
+    parser.add_argument('--featuresFolders', type=str, nargs='+', help='')
     parser.add_argument('--output', type=str, help='')
     parser.add_argument('--masksFolder', type=str, help='')
     parser.add_argument('--masksExtension', type=str, help='')
@@ -196,7 +206,7 @@ def compute_featuresMatcher(inputSfMData, imagePairsList, warpFolder, featuresFo
         args.func(inputSfMData=args.inputSfMData,
                 imagePairsList=args.imagePairsList,
                 warpFolder=args.warpFolder,
-                featuresFolder=args.featuresFolder,
+                featuresFolders=args.featuresFolders,
                 matchesFolder=args.output,
                 masksFolder=args.masksFolder,
                 masksExtension=args.masksExtension,

python/matcher.py

@@ -1,4 +1,7 @@
-from romatch import roma_outdoor
+from romav2.device import device
+from romav2 import RoMaV2
+from romav2.io import tensor_to_pil
+from romav2.features import Descriptor
 
 from common import *
 
@@ -8,32 +11,22 @@
 
 
 def prepare_warp(w):
-    """ Transform the warp tensor from roma to a RGB image with B value being always 1 """
-    w = ((w + 1.0) / 2.0).detach().cpu().numpy()
 
+    """ Transform the warp tensor from roma to a RGB image with B value being always 1 """
+    w = ((w + 1.0) / 2.0).detach().cpu().numpy().copy()
     w = np.concatenate([w, np.zeros([w.shape[0], w.shape[1], 1], dtype=np.float32)], axis=-1)
+
+
     return w
 
 def prepare_confidence(c):
     """ Transform the confidence tensor from roma to a 3 dimensional array 
     (Last dimension being of size 1)
     """
-    c = c.detach().cpu().numpy()
-    c = np.expand_dims(c, axis=-1)
+    c = c.detach().cpu().numpy().copy()
+
     return c
 
-def prepare_roma_outputs(w, c, upsampleResolution):
-    """ Transform output of roma to usable data
-    """
-    H = upsampleResolution[0]
-    W = upsampleResolution[1]
-
-    w_a_b = w[0, :, :W, 2:4]
-    c_a_b = c[0, :, :W]
-    w_b_a = w[0, :, W:, 0:2]
-    c_b_a = c[0, :, W:]
-
-    return prepare_warp(w_a_b), prepare_warp(w_b_a), prepare_confidence(c_a_b), prepare_confidence(c_b_a)
 
 def checkUncertaintyLoops(warp_A_B, warp_B_A, certainty_A_B, certainty_B_A, upsampleResolution):
     """ Take the minimum of certainty between the original certainty, and the certainty of the warped pixel.
@@ -80,7 +73,6 @@ def compute_densematches(inputSfMData, imagePairsList, outputWarpFolder, outputC
         outputCertaintyFolder : a destination folder for the certainty images
     """
 
-    upsampleResolution = (864, 864)
 
     #Parse sfmdata, create compatible images
     iinfos = get_imageinfos_from_sfmdata(inputSfMData)
@@ -102,16 +94,20 @@ def compute_densematches(inputSfMData, imagePairsList, outputWarpFolder, outputC
 
     logging.info("Loading model ....")
 
-    dinov2_weights = None
-    romaOutdoorModel = None
+    roma_weights = None
+    dinov3_path = None
     if "ROMATCH_MODELS_PATH" in os.environ:
         modelPath = os.environ["ROMATCH_MODELS_PATH"]
-        romaOutdoorModelPath = os.path.join(modelPath, "roma_outdoor.pth")
-        dinov2ModelPath = os.path.join(modelPath, "dinov2_vitl14_pretrain.pth")
-        dinov2_weights = torch.load(dinov2ModelPath, weights_only=True)
-        romaOutdoorModel = torch.load(romaOutdoorModelPath, weights_only=True)
-
-    matcher = roma_outdoor(device="cuda", upsample_res=upsampleResolution, weights=romaOutdoorModel, dinov2_weights=dinov2_weights)
+        romaModelPath = os.path.join(modelPath, "romav2.pt")
+        roma_weights = torch.load(romaModelPath, weights_only=True)
+        dinov3_path = os.path.join(modelPath, "dinov3")
+
+    descCfg = Descriptor.Cfg(module_path=dinov3_path)
+    romaCfg = RoMaV2.Cfg(descriptor=descCfg, weights=roma_weights)
+    model = RoMaV2(cfg=romaCfg)
+    model.apply_setting("precise")
+    upsampleResolution = (model.H_lr, model.W_lr) if (model.H_hr is None or model.W_hr is None) else (model.H_hr, model.W_hr) 
+
 
     for item in pairsToProcess:
         referenceId = item[0]
@@ -125,19 +121,23 @@ def compute_densematches(inputSfMData, imagePairsList, outputWarpFolder, outputC
 
         imA = open_image_to_pil(referenceInfo.path)
         imB = open_image_to_pil(otherInfo.path)
-        torch_warp, torch_certainty = matcher.match(imA, imB, device="cuda")
-
-        # prepare and stack data
-        warp_A_B, warp_B_A, certainty_A_B, certainty_B_A = prepare_roma_outputs(torch_warp, torch_certainty, upsampleResolution) 
 
+        preds = model.match(imA, imB)
+        warp_A_B = prepare_warp(preds["warp_AB"][0])
+        warp_B_A = prepare_warp(preds["warp_BA"][0])
+        certainty_A_B, certainty_B_A = (
+            prepare_confidence(preds["overlap_AB"][0]),
+            prepare_confidence(preds["overlap_BA"][0]),
+        )
+
         if checkLoops:
             checkUncertaintyLoops(warp_A_B, warp_B_A, certainty_A_B, certainty_B_A, upsampleResolution)
 
         logging.info("saving matches")
         pair_string = str(referenceId) + "_" + str(otherId)
         path_warp = os.path.join(outputWarpFolder, pair_string + "_warp.exr")
         path_certainty = os.path.join(outputCertaintyFolder, pair_string + "_certainty.exr")
-        save_image(path_warp, warp_A_B)
+        save_image(path_warp, warp_A_B, False)
         save_image(path_certainty, certainty_A_B, True)
 
 if __name__ == '__main__':