diff --git a/meshroom/roma/MatchMasking.py b/meshroom/roma/MatchMasking.py index c5c034a..5c876a6 100644 --- a/meshroom/roma/MatchMasking.py +++ b/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}" diff --git a/meshroom/roma/RomaFeaturesMatcher.py b/meshroom/roma/RomaFeaturesMatcher.py index e6f5540..0c098dd 100644 --- a/meshroom/roma/RomaFeaturesMatcher.py +++ b/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", diff --git a/meshroom/roma/RomaMatcher.py b/meshroom/roma/RomaMatcher.py index cd1dc66..29f38d1 100644 --- a/meshroom/roma/RomaMatcher.py +++ b/meshroom/roma/RomaMatcher.py @@ -2,8 +2,8 @@ import os from pathlib import Path - from meshroom.core import desc +from pyalicevision import parallelization as avpar class RomaMatcher(desc.CommandLineNode): @@ -11,7 +11,7 @@ class RomaMatcher(desc.CommandLineNode): 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) diff --git a/meshroom/roma/RomaReducer.py b/meshroom/roma/RomaReducer.py index 68d155e..98a2a50 100644 --- a/meshroom/roma/RomaReducer.py +++ b/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" ), ] diff --git a/meshroom/roma/RomaSampler.py b/meshroom/roma/RomaSampler.py index 96495ce..ee65e37 100644 --- a/meshroom/roma/RomaSampler.py +++ b/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" ), ] diff --git a/meshroom/roma/StarListing.py b/meshroom/roma/StarListing.py index 37a3ca0..5886a91 100644 --- a/meshroom/roma/StarListing.py +++ b/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( diff --git a/python/featuresMatcher.py b/python/featuresMatcher.py index c946ab1..ab975d6 100644 --- a/python/featuresMatcher.py +++ b/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, diff --git a/python/matcher.py b/python/matcher.py index fe76b4f..1cfcd4f 100644 --- a/python/matcher.py +++ b/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,11 +121,15 @@ 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) @@ -137,7 +137,7 @@ def compute_densematches(inputSfMData, imagePairsList, outputWarpFolder, outputC 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__':