Refactor skeleton.to_json to serialize skeleton object without jsonpickle

sleap/skeleton.py

@@ -987,9 +987,9 @@ 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:`Nodes` to index
                 in some list. This is used when saving
-                :class:`Labels`where we want to serialize the
+                :class:`Labels` where we want to serialize the
                 :class:`Nodes` outside the :class:`Skeleton` object.
                 If given, then we replace each :class:`Node` with
                 specified index before converting :class:`Skeleton`.
@@ -999,34 +999,114 @@ 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.
-        if self.is_template:
-            data = {
-                "nx_graph": graph,
-                "description": self.description,
-                "preview_image": self.preview_image,
-            }
-        else:
-            data = graph
+        # Logic taken from github.com/talmolab/sleap-io/io/slp.py::serialize_skeletons
+        # https://github.com/talmolab/sleap-io/blob/main/sleap_io/io/slp.py#L606
+        # Create a dictionary to store node data
+        nodes_dicts = []
+        node_to_id = {}
+        for node in self.nodes:
+            print(f'node: {node}')
+            if node not in node_to_id:
+                node_to_id[node] = node_to_idx[node] if node_to_idx is not None else len(node_to_id)
+                nodes_dicts.append({"name": node.name, "weight": 1.0})
+                print(f'node_to_id: {node_to_id}')
+                print(f'nodes_dicts: {nodes_dicts}')
+
+        # Create a dictionary to store edge data
+        edges_dicts = []
+        for edge_ind, edge in enumerate(self.edges):
+            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_dicts.append(
+                {
+                    "edge_insert_idx": edge_ind,
+                    "key": 0,  # Always 0.
+                    "source": node_to_id[edge[0]],
+                    "target": node_to_id[edge[1]],
+                    "type": edge_type,
+                }
+            )
+            print(f'edges_dicts: {edges_dicts}')
+        # Build links dicts for symmetry edges.
+        for symmetry_ind, symmetry in enumerate(self.symmetries):
+            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}
+
+            src, dst = tuple(symmetry.nodes)
+            edges_dicts.append(
+                {
+                    "key": 0,
+                    "source": node_to_id[src],
+                    "target": node_to_id[dst],
+                    "type": edge_type,
+                }
+            )
 
-        json_str = jsonpickle.encode(data)
+        # Create skeleton dict.
+        skeleton_dict = {
+            "directed": True,
+            "graph": {
+                "name": self.name,
+                "num_edges_inserted": len(self.edges),
+            },
+            "links": edges_dicts,
+            "multigraph": True,
+            # In the order in Skeleton.nodes and must match up with nodes_dicts.
+            "nodes": [{"id": node_to_id[node]} for node in self.nodes],
+        }
+
+        # Convert the skeleton dict to a JSON string using the standard json module
+        json_str = json.dumps(skeleton_dict, indent=4, sort_keys=True)
 
         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.