From d3ded2a32809bd7f41c8c4e86535d02687ed2e1e Mon Sep 17 00:00:00 2001
From: kijai <jukka.seppanen@poke.fi>
Date: Mon, 16 Oct 2023 18:46:46 +0300
Subject: [PATCH] Add linewidth as variable

Thicker maybe better for openpose cn
---
 __init__.py               | 5 +++--
 mogen/utils/plot_utils.py | 4 ++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/__init__.py b/__init__.py
index 375a1b9..a84b17e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -235,6 +235,7 @@ class MotionDataVisualizer:
                 "distance": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 10.0, "step": 0.1}),
                 "elevation": ("FLOAT", {"default": 120, "min": 0.0, "max": 300.0, "step": 0.1}),
                 "rotation": ("FLOAT", {"default": -90, "min": -180, "max": 180, "step": 1}),
+                "poselinewidth": ("FLOAT", {"default": 2, "min": 0, "max": 50, "step": 0.1}),
             },
             "optional": {
                 "opt_title": ("STRING", {"default": '' ,"multiline": False}),
@@ -245,12 +246,12 @@ class MotionDataVisualizer:
     CATEGORY = "MotionDiff"
     FUNCTION = "visualize"
 
-    def visualize(self, motion_data, visualization, distance, elevation, rotation, opt_title=None):
+    def visualize(self, motion_data, visualization, distance, elevation, rotation, poselinewidth, opt_title=None):
         pred_motion = motion_data["motion"]
         joint = recover_from_ric(pred_motion, 22).numpy()
         joint = motion_temporal_filter(joint, sigma=2.5)
         pil_frames = plot_3d_motion(
-            None, t2m_kinematic_chain, joint, distance, elevation, rotation,
+            None, t2m_kinematic_chain, joint, distance, elevation, rotation, poselinewidth,
             title=opt_title if opt_title is not None else '',
             fps=1,  save_as_pil_lists=True, visualization=visualization
         )
diff --git a/mogen/utils/plot_utils.py b/mogen/utils/plot_utils.py
index 55a857d..ca266bd 100644
--- a/mogen/utils/plot_utils.py
+++ b/mogen/utils/plot_utils.py
@@ -142,7 +142,7 @@ def recover_from_ric(data, joints_num):
 
     return positions
 
-def plot_3d_motion(save_path, kinematic_tree, joints, distance, elevation, rotation, title, figsize=(10, 10), fps=120,  radius=4, visualization="original", save_as_pil_lists=False):
+def plot_3d_motion(save_path, kinematic_tree, joints, distance, elevation, rotation, poselinewidth, title, figsize=(10, 10), fps=120,  radius=4, visualization="original", save_as_pil_lists=False):
     matplotlib.use('Agg')
 
     title_sp = title.split(' ')
@@ -226,7 +226,7 @@ def plot_3d_motion(save_path, kinematic_tree, joints, distance, elevation, rotat
                 linewidth = 2.0
             if visualization == "pseudo-openpose":
                 for j in range(len(data[index, chain, 0]) - 1):
-                  ax.plot3D(data[index, chain, 0][j:j+2], data[index, chain, 1][j:j+2], data[index, chain, 2][j:j+2], linewidth=linewidth,
+                  ax.plot3D(data[index, chain, 0][j:j+2], data[index, chain, 1][j:j+2], data[index, chain, 2][j:j+2], linewidth=poselinewidth,
                             color=color[j % len(color)])
                   
             else: