Refactor skeleton.to_json to serialize skeleton object without jsonpickle

sleap/skeleton.py

@@ -987,7 +987,7 @@ def to_json(self, node_to_idx: Optional[Dict[Node, int]] = None) -> str:
         """Convert the :class:`Skeleton` to a JSON representation.
 
         Args:
-            node_to_idx: optional dict which maps :class:`Node`sto index
+            node_to_idx: optional dict which maps :class:`Node`s to index
                 in some list. This is used when saving
                 :class:`Labels`where we want to serialize the
                 :class:`Nodes` outside the :class:`Skeleton` object.
@@ -999,34 +999,167 @@ def to_json(self, node_to_idx: Optional[Dict[Node, int]] = None) -> str:
         Returns:
             A string containing the JSON representation of the skeleton.
         """
-        jsonpickle.set_encoder_options("simplejson", sort_keys=True, indent=4)
-        if node_to_idx is not None:
-            indexed_node_graph = nx.relabel_nodes(
-                G=self._graph, mapping=node_to_idx
-            )  # map nodes to int
-        else:
-            indexed_node_graph = self._graph
-
-        # Encode to JSON
-        graph = json_graph.node_link_data(indexed_node_graph)
-
-        # SLEAP v1.3.0 added `description` and `preview_image` to `Skeleton`, but saving
-        # these fields breaks data format compatibility. Currently, these are only
-        # added in our custom template skeletons. To ensure backwards data format
-        # compatibilty of user data, we only save these fields if they are not None.
+        # Create global list of nodes with all nodes from all skeletons.
+        nodes_dicts = []
+        node_to_id = {}
+        for node in self.nodes:
+            if node not in node_to_id:
+                print(f"node: {node}")
+                # Note: This ID is not the same as the node index in the skeleton in
+                # legacy SLEAP, but we do not retain this information in the labels, so
+                # IDs will be different.
+                #
+                # The weight is also kept fixed here, but technically this is not
+                # modified or used in legacy SLEAP either.
+                #
+                # TODO: Store legacy metadata in labels to get byte-level compatibility?
+                node_to_id[node] = len(node_to_id)
+                print(f"node_to_id: {node_to_id}")
+                nodes_dicts.append({"name": node.name, "weight": 1.0})
+        print(f"nodes_dicts: {nodes_dicts}")
+
+        # Build links dicts for normal edges.
+        edges_dicts = []
+        for edge_ind, edge in enumerate(self.edges):
+            print(f"edge_ind: {edge_ind}")
+            print(f"edge: {edge}")
+            if edge_ind == 0:
+                edge_type = {
+                    "py/reduce": [
+                        {"py/type": "sleap.skeleton.EdgeType"},
+                        {"py/tuple": [1]},  # 1 = real edge, 2 = symmetry edge
+                    ]
+                }
+                print(f"edge_type: {edge_type}")
+            else:
+                edge_type = {"py/id": 1}
+                print(f"edge_type: {edge_type}")
+
+            # Edges are stored as a list of tuples of nodes
+            # The source and target are the nodes in the tuple (edge) are the first and
+            # second nodes respectively
+            source = edge[0]
+            print(f"source: {source}")
+            print(f"node_to_id[source]: {node_to_id[source]}")
+            target = edge[1]
+            print(f"target: {target}")
+            print(f"node_to_id[target]: {node_to_id[target]}")
+            edges_dicts.append(
+                {
+                    # Note: Insert idx is not the same as the edge index in the skeleton
+                    # in legacy SLEAP.
+                    "edge_insert_idx": edge_ind,
+                    "key": 0,  # Always 0.
+                    "source": {"py/id": self.nodes.index(node)},
+                    "target": {"py/id": self.nodes.index(node)},
+                    "type": edge_type,
+                }
+            )
+        print(f"edges_dicts: {edges_dicts}")
+
+        # Build links dicts for symmetry edges.
+        for symmetry_ind, symmetry in enumerate(self.symmetries):
+            print(f"symmetry_ind: {symmetry_ind}")
+            print(f"symmetry: {symmetry}")
+            if symmetry_ind == 0:
+                edge_type = {
+                    "py/reduce": [
+                        {"py/type": "sleap.skeleton.EdgeType"},
+                        {"py/tuple": [2]},  # 1 = real edge, 2 = symmetry edge
+                    ]
+                }
+            else:
+                edge_type = {"py/id": 2}
+
+            symmetry_src, symmetry_dst = symmetry
+            print(f"src: {symmetry_src}")
+            print(f"dst: {symmetry_dst}")
+            edges_dicts.append(
+                {
+                    "key": 0,
+                    "source": {"py/id": node_to_id[symmetry_src]},
+                    "target": {"py/id": node_to_id[symmetry_dst]},
+                    "type": edge_type,
+                }
+            )
+        # Create graph field
+        graph = {
+            "name": self.name,
+            "num_edges_inserted": len(self.edges),
+        }
+        # Create skeleton dict
         if self.is_template:
-            data = {
-                "nx_graph": graph,
+            # Template skeletons have additional fields
+            nx_graph = {
+                "directed": True,
+                "graph": graph,
+                "nodes": [
+                    {
+                        "id": {
+                            "py/object": "sleap.skeleton.Node",
+                            "py/state": {"name": node.name, "weight": node.weight},
+                        }
+                    }
+                    for node in self.nodes
+                ],
+                "links": edges_dicts,
+                "multigraph": True,
+            }
+            skeleton_dict = {
                 "description": self.description,
+                "nx_graph": nx_graph,
                 "preview_image": self.preview_image,
             }
         else:
-            data = graph
-
-        json_str = jsonpickle.encode(data)
+            skeleton_dict = {
+                "directed": True,
+                "graph": graph,
+                "nodes": [
+                    {
+                        "id": {
+                            "py/object": "sleap.skeleton.Node",
+                            "py/state": {"name": node.name, "weight": node.weight},
+                        }
+                    }
+                    for node in self.nodes
+                ],
+                "links": edges_dicts,
+                "multigraph": True,
+            }
 
+        print(f"skeleton_dict: {skeleton_dict}")
+        json_str = json.dumps(skeleton_dict, indent=4, sort_keys=True)
+        print(f"json_str: {json_str}")
         return json_str
 
+        # jsonpickle.set_encoder_options("simplejson", sort_keys=True, indent=4)
+        # if node_to_idx is not None:
+        #     indexed_node_graph = nx.relabel_nodes(
+        #         G=self._graph, mapping=node_to_idx
+        #     )  # map nodes to int
+        # else:
+        #     indexed_node_graph = self._graph
+
+        # # Encode to JSON
+        # graph = json_graph.node_link_data(indexed_node_graph)
+
+        # # SLEAP v1.3.0 added `description` and `preview_image` to `Skeleton`, but saving
+        # # these fields breaks data format compatibility. Currently, these are only
+        # # added in our custom template skeletons. To ensure backwards data format
+        # # compatibilty of user data, we only save these fields if they are not None.
+        # if self.is_template:
+        #     data = {
+        #         "nx_graph": graph,
+        #         "description": self.description,
+        #         "preview_image": self.preview_image,
+        #     }
+        # else:
+        #     data = graph
+
+        # json_str = jsonpickle.encode(data)
+
+        # return json_str
+
     def save_json(self, filename: str, node_to_idx: Optional[Dict[Node, int]] = None):
         """
         Save the :class:`Skeleton` as JSON file.