diff --git a/bin/meshroom_batch b/bin/meshroom_batch index 7883e1b6f8..c5363acd52 100755 --- a/bin/meshroom_batch +++ b/bin/meshroom_batch @@ -9,8 +9,10 @@ meshroom.setupEnvironment() import meshroom.core.graph from meshroom import multiview +from meshroom.core.desc import InitNode import logging + parser = argparse.ArgumentParser(description='Launch the full photogrammetry or Panorama HDR pipeline.') parser.add_argument('-i', '--input', metavar='SFM/FOLDERS/IMAGES', type=str, nargs='*', default=[], @@ -94,33 +96,24 @@ def getOnlyNodeOfType(g, nodeType): return nodes[0] +def getInitNode(g): + """ + Helper function to get the Init node in the graph 'g' and raise an exception if there is no or + multiple candidates. + """ + nodes = g.findInitNodes() + if len(nodes) == 0: + raise RuntimeError("meshroom_batch requires an Init node in the pipeline.") + elif len(nodes) > 1: + raise RuntimeError("meshroom_batch requires exactly one Init node in the pipeline, {} found: {}" + .format(len(nodes), str(nodes))) + return nodes[0] + + if not args.input and not args.inputRecursive: print('Nothing to compute. You need to set --input or --inputRecursive.') sys.exit(1) -views, intrinsics = [], [] -# Build image files list from inputImages arguments -filesByType = multiview.FilesByType() - -hasSearchedForImages = False - -if args.input: - if len(args.input) == 1 and os.path.isfile(args.input[0]) and os.path.splitext(args.input[0])[-1] in ('.json', '.sfm'): - # args.input is a sfmData file: setup pre-calibrated views and intrinsics - from meshroom.nodes.aliceVision.CameraInit import readSfMData - views, intrinsics = readSfMData(args.input[0]) - else: - filesByType.extend(multiview.findFilesByTypeInFolder(args.input, recursive=False)) - hasSearchedForImages = True - -if args.inputRecursive: - filesByType.extend(multiview.findFilesByTypeInFolder(args.inputRecursive, recursive=True)) - hasSearchedForImages = True - -if hasSearchedForImages and not filesByType.images: - print("No image found") - sys.exit(-1) - graph = multiview.Graph(name=args.pipeline) with multiview.GraphModification(graph): @@ -131,15 +124,10 @@ with multiview.GraphModification(graph): else: # custom pipeline graph.load(args.pipeline, setupProjectFile=False) - # graph.update() - cameraInit = getOnlyNodeOfType(graph, 'CameraInit') - # reset graph inputs - cameraInit.viewpoints.resetValue() - cameraInit.intrinsics.resetValue() - # add views and intrinsics (if any) read from args.input - cameraInit.viewpoints.extend(views) - cameraInit.intrinsics.extend(intrinsics) + # get init node and initialize it + initNode = getInitNode(graph) + initNode.nodeDesc.initialize(initNode, args.input, args.inputRecursive) if not graph.canComputeLeaves: raise RuntimeError("Graph cannot be computed. Check for compatibility issues.") @@ -151,11 +139,6 @@ with multiview.GraphModification(graph): publish = getOnlyNodeOfType(graph, 'Publish') publish.output.value = args.output - if filesByType.images: - views, intrinsics = cameraInit.nodeDesc.buildIntrinsics(cameraInit, filesByType.images) - cameraInit.viewpoints.value = views - cameraInit.intrinsics.value = intrinsics - if args.overrides: import io import json diff --git a/meshroom/core/desc.py b/meshroom/core/desc.py index 3dcb7a0748..e66eb053a7 100755 --- a/meshroom/core/desc.py +++ b/meshroom/core/desc.py @@ -527,3 +527,55 @@ def processChunk(self, chunk): finally: chunk.subprocess = None + +# Test abstract node +class InitNode: + def __init__(self): + pass + + def initialize(self, node, inputs, recursiveInputs): + """ + Initialize the attributes that are needed for a node to start running. + + Args: + node (Node): the node whose attributes must be initialized + inputs (list): the user-provided list of input files/directories + recursiveInputs (list): the user-provided list of input directories to search recursively for images + """ + pass + + def resetAttributes(self, node, attributeNames): + """ + Reset the values of the provided attributes for a node. + + Args: + node (Node): the node whose attributes are to be reset + attributeNames (list): the list containing the names of the attributes to reset + """ + for attrName in attributeNames: + if node.hasAttribute(attrName): + node.attribute(attrName).resetValue() + + def extendAttributes(self, node, attributesDict): + """ + Extend the values of the provided attributes for a node. + + Args: + node (Node): the node whose attributes are to be extended + attributesDict (dict): the dictionary containing the attributes' names (as keys) and the values to extend with + """ + for attr in attributesDict.keys(): + if node.hasAttribute(attr): + node.attribute(attr).extend(attributesDict[attr]) + + def setAttributes(self, node, attributesDict): + """ + Set the values of the provided attributes for a node. + + Args: + node (Node): the node whose attributes are to be extended + attributesDict (dict): the dictionary containing the attributes' names (as keys) and the values to set + """ + for attr in attributesDict: + if node.hasAttribute(attr): + node.attribute(attr).value = attributesDict[attr] diff --git a/meshroom/core/graph.py b/meshroom/core/graph.py index 1e9b81329a..89faefb51d 100644 --- a/meshroom/core/graph.py +++ b/meshroom/core/graph.py @@ -547,6 +547,14 @@ def nodesOfType(self, nodeType, sortedByIndex=True): nodes = [n for n in self._nodes.values() if n.nodeType == nodeType] return self.sortNodesByIndex(nodes) if sortedByIndex else nodes + def findInitNodes(self): + """ + Returns: + list[Node]: the list of Init nodes (nodes inheriting from InitNode) + """ + nodes = [n for n in self._nodes.values() if isinstance(n.nodeDesc, meshroom.core.desc.InitNode)] + return nodes + def findNodeCandidates(self, nodeNameExpr): pattern = re.compile(nodeNameExpr) return [v for k, v in self._nodes.objects.items() if pattern.match(k)] diff --git a/meshroom/nodes/aliceVision/CameraInit.py b/meshroom/nodes/aliceVision/CameraInit.py index 9ff8a83457..dcc8182b6c 100644 --- a/meshroom/nodes/aliceVision/CameraInit.py +++ b/meshroom/nodes/aliceVision/CameraInit.py @@ -8,7 +8,7 @@ import logging from meshroom.core import desc, Version - +from meshroom.multiview import FilesByType, findFilesByTypeInFolder Viewpoint = [ desc.IntParam(name="viewId", label="Id", description="Image UID", value=-1, uid=[0], range=None), @@ -119,7 +119,8 @@ def readSfMData(sfmFile): return views, intrinsics -class CameraInit(desc.CommandLineNode): + +class CameraInit(desc.CommandLineNode, desc.InitNode): commandLine = 'aliceVision_cameraInit {allParams} --allowSingleView 1' # don't throw an error if there is only one image size = desc.DynamicNodeSize('viewpoints') @@ -250,6 +251,35 @@ class CameraInit(desc.CommandLineNode): ), ] + def __init__(self): + super(CameraInit, self).__init__() + + def initialize(self, node, inputs, recursiveInputs): + # Reset graph inputs + self.resetAttributes(node, ["viewpoints", "intrinsics"]) + + filesByType = FilesByType() + searchedForImages = False + + if recursiveInputs: + filesByType.extend(findFilesByTypeInFolder(recursiveInputs, recursive=True)) + searchedForImages = True + + # Add views and intrinsics from a file if it was provided, or look for the images + if len(inputs) == 1 and os.path.isfile(inputs[0]) and os.path.splitext(inputs[0])[-1] in ('.json', '.sfm'): + views, intrinsics = readSfMData(inputs[0]) + self.extendAttributes(node, {"viewpoints": views, "intrinsics": intrinsics}) + else: + filesByType.extend(findFilesByTypeInFolder(inputs, recursive=False)) + searchedForImages = True + + # If there was no input file, check that the directories do contain images + if searchedForImages and not filesByType.images: + raise ValueError("No valid input file or no image in the provided directories") + + views, intrinsics = self.buildIntrinsics(node, filesByType.images) + self.setAttributes(node, {"viewpoints": views, "intrinsics": intrinsics}) + def upgradeAttributeValues(self, attrValues, fromVersion): # Starting with version 6, the principal point is now relative to the image center