Merge pull request #1750 from alicevision/dev/abstractCamInit

Add abstract InitNode

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

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]

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)]

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