From 096f04d164bf2ea6ab2c1c2c7469f3c4e5aedf3d Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 15 Dec 2025 00:52:25 +0200
Subject: [PATCH] Draw 2nd person in correct colors

---
 NLFPoseExtract/nlf_render.py | 18 +++++++++++++-----
 nodes.py                     |  7 +++++--
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/NLFPoseExtract/nlf_render.py b/NLFPoseExtract/nlf_render.py
index b0fe7e6..4b0c11e 100644
--- a/NLFPoseExtract/nlf_render.py
+++ b/NLFPoseExtract/nlf_render.py
@@ -256,10 +256,8 @@ def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intr
     return frames_np_rgba
 
 
-def render_multi_nlf_as_images(data, dw_poses, intrinsic_matrix=None, draw_2d=True):
-    """ return a list of images """
-    height, width = data[0]['video_height'], data[0]['video_width']
-    video_length = len(data)
+def render_multi_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True, draw_face=True, draw_hands=True):
+
 
     second_person_base_colors_255_dict = {
         # Warm Colors for Right Side (R.) - Red, Orange, Yellow
@@ -368,8 +366,18 @@ def render_multi_nlf_as_images(data, dw_poses, intrinsic_matrix=None, draw_2d=Tr
     colors_first = [[c / 300 + 0.15 for c in color_rgb] + [0.8] for color_rgb in ordered_colors_255_list[0]]
     colors_second = [[c / 300 + 0.15 for c in color_rgb] + [0.8] for color_rgb in ordered_colors_255_list[1]]
 
-    smpl_poses_first, smpl_poses_second = collect_smpl_poses_samurai(data)
+    smpl_poses_first = []
+    smpl_poses_second = []
+    for i in range(video_length):
+        if len(smpl_poses[i]) >= 1:
+            smpl_poses_first.append([smpl_poses[i][0]])  # First person
+        else:
+            smpl_poses_first.append([torch.zeros((24, 3), dtype=torch.float32)])
 
+        if len(smpl_poses[i]) >= 2:
+            smpl_poses_second.append([smpl_poses[i][1]])  # Second person
+        else:
+            smpl_poses_second.append([torch.zeros((24, 3), dtype=torch.float32)])
 
     if intrinsic_matrix is None:
         intrinsic_matrix = intrinsic_matrix_from_field_of_view((height, width))
diff --git a/nodes.py b/nodes.py
index 4aa2cb0..bd8f169 100644
--- a/nodes.py
+++ b/nodes.py
@@ -176,7 +176,7 @@ class RenderNLFPoses:
 
     def predict(self, nlf_poses, width, height, dw_poses=None, ref_dw_pose=None, draw_face=True, draw_hands=True):
 
-        from .NLFPoseExtract.nlf_render import render_nlf_as_images, shift_dwpose_according_to_nlf, process_data_to_COCO_format, intrinsic_matrix_from_field_of_view
+        from .NLFPoseExtract.nlf_render import render_nlf_as_images, render_multi_nlf_as_images, shift_dwpose_according_to_nlf, process_data_to_COCO_format, intrinsic_matrix_from_field_of_view
         from .NLFPoseExtract.align3d import solve_new_camera_params_central, solve_new_camera_params_down
 
         if isinstance(nlf_poses, dict):
@@ -230,7 +230,10 @@ class RenderNLFPoses:
         else:
             intrinsic_matrix = ori_camera_pose
 
-        frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands)
+        if pose_input[0].shape[0] > 1:
+            frames_np = render_multi_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands)
+        else:
+            frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands)
 
         frames_tensor = torch.from_numpy(np.stack(frames_np, axis=0)).contiguous() / 255.0
         frames_tensor, mask = frames_tensor[..., :3], frames_tensor[..., -1] > 0.5