fix primMST_symmetry bug

VietAnhDao · Mar 30, 2021 · 24ca71b · 24ca71b
1 parent e5f6ab7
commit 24ca71b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/mst_generate.py b/mst_generate.py
@@ -202,7 +202,7 @@ def run_mst_generate(args):
         cost_matrix = increase_cost_for_outside_bone(cost_matrix, pred_joints, vox)
 
         skel = Skel()
-        parent, key = primMST_symmetry(cost_matrix, root_id, pred_joints)
+        parent, key, root_id = primMST_symmetry(cost_matrix, root_id, pred_joints)
         for i in range(len(parent)):
             if parent[i] == -1:
                 skel.root = TreeNode('root', tuple(pred_joints[i]))

diff --git a/quick_start.py b/quick_start.py
@@ -201,7 +201,7 @@ def predict_skeleton(input_data, vox, root_pred_net, bone_pred_net, mesh_filenam
     cost_matrix = increase_cost_for_outside_bone(cost_matrix, pred_joints, vox)
 
     pred_skel = Info()
-    parent, key = primMST_symmetry(cost_matrix, root_id, pred_joints)
+    parent, key, root_id = primMST_symmetry(cost_matrix, root_id, pred_joints)
     for i in range(len(parent)):
         if parent[i] == -1:
             pred_skel.root = TreeNode('root', tuple(pred_joints[i]))

diff --git a/utils/mst_utils.py b/utils/mst_utils.py
@@ -185,7 +185,7 @@ def primMST_symmetry(graph, init_id, joints):
                 key[v] = graph[u2, v]
                 parent[v] = u2
 
-    return parent, key
+    return parent, key, init_id
 
 
 def loadSkel_recur(p_node, parent_id, joint_name, joint_pos, parent):