Merge pull request #2612 from alicevision/dev/graphIO

Refactor Graph de/serialization

Author: cbentejac

bin/meshroom_batch

@@ -154,10 +154,10 @@ with meshroom.core.graph.GraphModification(graph):
     # initialize template pipeline
     loweredPipelineTemplates = dict((k.lower(), v) for k, v in meshroom.core.pipelineTemplates.items())
     if args.pipeline.lower() in loweredPipelineTemplates:
-        graph.load(loweredPipelineTemplates[args.pipeline.lower()], setupProjectFile=False, publishOutputs=True if args.output else False)
+        graph.initFromTemplate(loweredPipelineTemplates[args.pipeline.lower()], publishOutputs=True if args.output else False)
     else:
         # custom pipeline
-        graph.load(args.pipeline, setupProjectFile=False, publishOutputs=True if args.output else False)
+        graph.initFromTemplate(args.pipeline, publishOutputs=True if args.output else False)
 
     def parseInputs(inputs, uniqueInitNode):
         """Utility method for parsing the input and inputRecursive arguments."""

meshroom/core/attribute.py

@@ -339,9 +339,12 @@ def _applyExpr(self):
         elif self.isInput and Attribute.isLinkExpression(v):
             # value is a link to another attribute
             link = v[1:-1]
-            linkNode, linkAttr = link.split('.')
+            linkNodeName, linkAttrName = link.split('.')
             try:
-                g.addEdge(g.node(linkNode).attribute(linkAttr), self)
+                node = g.node(linkNodeName)
+                if not node:
+                    raise KeyError(f"Node '{linkNodeName}' not found")
+                g.addEdge(node.attribute(linkAttrName), self)
             except KeyError as err:
                 logging.warning('Connect Attribute from Expression failed.')
                 logging.warning('Expression: "{exp}"\nError: "{err}".'.format(exp=v, err=err))

meshroom/core/graph.py

@@ -0,0 +1,231 @@
+from enum import Enum
+from typing import Any, TYPE_CHECKING, Union
+
+import meshroom
+from meshroom.core import Version
+from meshroom.core.attribute import Attribute, GroupAttribute, ListAttribute
+from meshroom.core.node import Node
+
+if TYPE_CHECKING:
+    from meshroom.core.graph import Graph
+
+
+class GraphIO:
+    """Centralize Graph file keys and IO version."""
+
+    __version__ = "2.0"
+
+    class Keys(object):
+        """File Keys."""
+
+        # Doesn't inherit enum to simplify usage (GraphIO.Keys.XX, without .value)
+        Header = "header"
+        NodesVersions = "nodesVersions"
+        ReleaseVersion = "releaseVersion"
+        FileVersion = "fileVersion"
+        Graph = "graph"
+        Template = "template"
+
+    class Features(Enum):
+        """File Features."""
+
+        Graph = "graph"
+        Header = "header"
+        NodesVersions = "nodesVersions"
+        PrecomputedOutputs = "precomputedOutputs"
+        NodesPositions = "nodesPositions"
+
+    @staticmethod
+    def getFeaturesForVersion(fileVersion: Union[str, Version]) -> tuple["GraphIO.Features", ...]:
+        """Return the list of supported features based on a file version.
+
+        Args:
+            fileVersion (str, Version): the file version
+
+        Returns:
+            tuple of GraphIO.Features: the list of supported features
+        """
+        if isinstance(fileVersion, str):
+            fileVersion = Version(fileVersion)
+
+        features = [GraphIO.Features.Graph]
+        if fileVersion >= Version("1.0"):
+            features += [
+                GraphIO.Features.Header,
+                GraphIO.Features.NodesVersions,
+                GraphIO.Features.PrecomputedOutputs,
+            ]
+
+        if fileVersion >= Version("1.1"):
+            features += [GraphIO.Features.NodesPositions]
+
+        return tuple(features)
+
+
+class GraphSerializer:
+    """Standard Graph serializer."""
+
+    def __init__(self, graph: "Graph") -> None:
+        self._graph = graph
+
+    def serialize(self) -> dict:
+        """
+        Serialize the Graph.
+        """
+        return {
+            GraphIO.Keys.Header: self.serializeHeader(),
+            GraphIO.Keys.Graph: self.serializeContent(),
+        }
+
+    @property
+    def nodes(self) -> list[Node]:
+        return self._graph.nodes
+
+    def serializeHeader(self) -> dict:
+        """Build and return the graph serialization header.
+
+        The header contains metadata about the graph, such as the:
+            - version of the software used to create it.
+            - version of the file format.
+            - version of the nodes types used in the graph.
+            - template flag.
+        """
+        header: dict[str, Any] = {}
+        header[GraphIO.Keys.ReleaseVersion] = meshroom.__version__
+        header[GraphIO.Keys.FileVersion] = GraphIO.__version__
+        header[GraphIO.Keys.NodesVersions] = self._getNodeTypesVersions()
+        return header
+
+    def _getNodeTypesVersions(self) -> dict[str, str]:
+        """Get registered versions of each node types in `nodes`, excluding CompatibilityNode instances."""
+        nodeTypes = set([node.nodeDesc.__class__ for node in self.nodes if isinstance(node, Node)])
+        nodeTypesVersions = {
+            nodeType.__name__: version
+            for nodeType in nodeTypes
+            if (version := meshroom.core.nodeVersion(nodeType)) is not None
+        }
+        # Sort them by name (to avoid random order changing from one save to another).
+        return dict(sorted(nodeTypesVersions.items()))
+
+    def serializeContent(self) -> dict:
+        """Graph content serialization logic."""
+        return {node.name: self.serializeNode(node) for node in sorted(self.nodes, key=lambda n: n.name)}
+
+    def serializeNode(self, node: Node) -> dict:
+        """Node serialization logic."""
+        return node.toDict()
+
+
+class TemplateGraphSerializer(GraphSerializer):
+    """Serializer for serializing a graph as a template."""
+
+    def serializeHeader(self) -> dict:
+        header = super().serializeHeader()
+        header[GraphIO.Keys.Template] = True
+        return header
+
+    def serializeNode(self, node: Node) -> dict:
+        """Adapt node serialization to template graphs.
+        
+        Instead of getting all the inputs and internal attribute keys, only get the keys of
+        the attributes whose value is not the default one.
+        The output attributes, UIDs, parallelization parameters and internal folder are
+        not relevant for templates, so they are explicitly removed from the returned dictionary.
+        """
+        # For now, implemented as a post-process to update the default serialization.
+        nodeData = super().serializeNode(node)
+
+        inputKeys = list(nodeData["inputs"].keys())
+
+        internalInputKeys = []
+        internalInputs = nodeData.get("internalInputs", None)
+        if internalInputs:
+            internalInputKeys = list(internalInputs.keys())
+
+        for attrName in inputKeys:
+            attribute = node.attribute(attrName)
+            # check that attribute is not a link for choice attributes
+            if attribute.isDefault and not attribute.isLink:
+                del nodeData["inputs"][attrName]
+
+        for attrName in internalInputKeys:
+            attribute = node.internalAttribute(attrName)
+            # check that internal attribute is not a link for choice attributes
+            if attribute.isDefault and not attribute.isLink:
+                del nodeData["internalInputs"][attrName]
+
+        # If all the internal attributes are set to their default values, remove the entry
+        if len(nodeData["internalInputs"]) == 0:
+            del nodeData["internalInputs"]
+
+        del nodeData["outputs"]
+        del nodeData["uid"]
+        del nodeData["internalFolder"]
+        del nodeData["parallelization"]
+
+        return nodeData
+
+
+class PartialGraphSerializer(GraphSerializer):
+    """Serializer to serialize a partial graph (a subset of nodes)."""
+
+    def __init__(self, graph: "Graph", nodes: list[Node]):
+        super().__init__(graph)
+        self._nodes = nodes
+
+    @property
+    def nodes(self) -> list[Node]:
+        """Override to consider only the subset of nodes."""
+        return self._nodes
+
+    def serializeNode(self, node: Node) -> dict:
+        """Adapt node serialization to partial graph serialization."""
+        # NOTE: For now, implemented as a post-process to the default serialization.
+        nodeData = super().serializeNode(node)
+
+        # Override input attributes with custom serialization logic, to handle attributes
+        # connected to nodes that are not in the list of nodes to serialize.
+        for attributeName in nodeData["inputs"]:
+            nodeData["inputs"][attributeName] = self._serializeAttribute(node.attribute(attributeName))
+
+        # Clear UID for non-compatibility nodes, as the custom attribute serialization
+        # can be impacting the UID by removing connections to missing nodes.
+        if not node.isCompatibilityNode:
+            del nodeData["uid"]
+
+        return nodeData
+
+    def _serializeAttribute(self, attribute: Attribute) -> Any:
+        """
+        Serialize `attribute` (recursively for list/groups) and deal with attributes being connected
+        to nodes that are not part of the partial list of nodes to serialize.
+        """
+        linkParam = attribute.getLinkParam()
+
+        if linkParam is not None:
+            # Use standard link serialization if upstream node is part of the serialization.
+            if linkParam.node in self.nodes:
+                return attribute.getExportValue()
+            # Skip link serialization otherwise.
+            # If part of a list, this entry can be discarded.
+            if isinstance(attribute.root, ListAttribute):
+                return None
+            # Otherwise, return the default value for this attribute.
+            return attribute.defaultValue()
+
+        if isinstance(attribute, ListAttribute):
+            # Recusively serialize each child of the ListAttribute, skipping those for which the attribute
+            # serialization logic above returns None.
+            return [
+                exportValue
+                for child in attribute
+                if (exportValue := self._serializeAttribute(child)) is not None
+            ]
+
+        if isinstance(attribute, GroupAttribute):
+            # Recursively serialize each child of the group attribute.
+            return {name: self._serializeAttribute(child) for name, child in attribute.value.items()}
+
+        return attribute.getExportValue()
+
+