Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactor skeleton.to_json to serialize skeleton object without jsonpickle #1934

Closed
wants to merge 20 commits into from
Closed

Refactor skeleton.to_json to serialize skeleton object without jsonpickle #1934

Changes from 15 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
cb62a2d
comment in `to_json` to start
eberrigan Sep 3, 2024
4c03ca2
fix confusing typo
eberrigan Sep 3, 2024
f352d03
add logic from sleap-io `serialize_skeletons` to `Skeleton.to_json` a…
eberrigan Sep 4, 2024
94dccb0
clean docstring
eberrigan Sep 4, 2024
4febed0
replace relabel nodes logic using networkx with custom indexing logic
eberrigan Sep 4, 2024
c3448f4
make dictionary directly instead of a list with a dictionary and save…
eberrigan Sep 4, 2024
3a7e65a
no attributes edge.source and edge.destination in SLEAP
eberrigan Sep 5, 2024
08b43ed
Merge branch 'develop' into elizabeth/refactor-skeletons-serialization
eberrigan Sep 10, 2024
5decb55
add to_json functions using sleap-io logic
eberrigan Sep 10, 2024
622853c
save changes to serialize skeletons
eberrigan Sep 11, 2024
884da47
match everything except indexing to legacy sleap
eberrigan Sep 11, 2024
06f18ab
Merge branch 'develop' into elizabeth/refactor-skeletons-serialization
eberrigan Sep 12, 2024
6e37ca9
move nodes list before links
eberrigan Sep 13, 2024
01d0255
format
eberrigan Sep 13, 2024
71b0307
symmetry is a tuple of nodes
eberrigan Sep 13, 2024
69db671
source and target should be tuples in edge dict
eberrigan Sep 13, 2024
0e2f9b5
check-in notebook for testing
eberrigan Sep 17, 2024
4ff6fdb
refer to source and target nodes using "py/id"
eberrigan Sep 17, 2024
f0451c5
use source and target nodes to index
eberrigan Sep 17, 2024
b4f10af
use tuple format for Node
eberrigan Sep 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
182 changes: 161 additions & 21 deletions sleap/skeleton.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +990 to 991
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove unused parameter node_to_idx from to_json method

The parameter node_to_idx is not used within the to_json method. Consider removing it to simplify the method signature and avoid confusion.

:class:`Labels`where we want to serialize the
:class:`Nodes` outside the :class:`Skeleton` object.
Expand All @@ -999,34 +999,174 @@ 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}")
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove debug print statements

The print statements throughout the to_json method appear to be for debugging purposes. These should be removed or replaced with appropriate logging to maintain clean production code.

Apply this diff to remove the print statements:

- print(f"node: {node}")
- print(f"node_to_id: {node_to_id}")
- print(f"nodes_dicts: {nodes_dicts}")
- print(f"edge_ind: {edge_ind}")
- print(f"edge: {edge}")
- print(f"edge_type: {edge_type}")
- print(f"source: {source}")
- print(f"node_to_id[source]: {node_to_id[source]}")
- print(f"target: {target}")
- print(f"node_to_id[target]: {node_to_id[target]}")
- print(f"edges_dicts: {edges_dicts}")
- print(f"symmetry_ind: {symmetry_ind}")
- print(f"symmetry: {symmetry}")
- print(f"src: {symmetry_src}")
- print(f"dst: {symmetry_dst}")
- print(f"skeleton_dict: {skeleton_dict}")
- print(f"json_str: {json_str}")

Also applies to: 1017-1017, 1019-1019, 1024-1026, 1033-1036, 1042-1046, 1058-1058, 1062-1063, 1075-1076, 1130-1132

# 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)
talmo marked this conversation as resolved.
Show resolved Hide resolved
Comment on lines +1015 to +1016
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Acknowledged the TODO for storing legacy metadata.

The TODO comment is a good reminder for a future enhancement to improve compatibility with legacy data.

Let me know if you need any help with implementing the storage of legacy metadata in labels to achieve byte-level compatibility. I can assist with the implementation or open a GitHub issue to track this task.

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/object": "sleap.skeleton.Node",
"py/state": {"name": source.name, "weight": 1.0},
},
"target": {
"py/object": "sleap.skeleton.Node",
"py/state": {"name": target.name, "weight": 1.0},
},
# "target": {"py/id": node_to_id[target]},
"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.
Expand Down
Loading