2.7.4 人体姿态估计

  """
用 MediaPipe 库的人体姿态估计提取图像中人体关键点信息，包括关键点之间的连线。
Input:
    ImgPath: Item Access; str
        图像文件路径
    StaticImageMode: Item Access; bool
        如果设置为 False，则将输入的多个图像视为视频流（video stream）。。如果设置为 True，则将输入的多个图像视为静态，可能并不相关的图像，对每个图像均执行人体检测
    ModelComplexity: Item Access; int
        配置模型的复杂度，其分为0，1 或 2 三个级别，默认为 1。值越大预测的关键点精度越高，但推断时长也会随之增加
    MinDetectionConfidence: Item Access; float
        人体检测模型（person-detection model）最小置信度，值域为 [0.0,1.0]
    HeightRef: Item Access; float
        参考位置高度
    kneeAngleLength: Item Access; float
        由小腿长度，用于控制人体身高
Output:
    X: List[float]
        关键点坐标 X 值
    Y: List[float]
        关键点坐标 Y 值    
    Z: List[float]
        关键点坐标 Z 值  
    XYZ: Tree[float]
        关键点坐标 XYZ 次序
    XZY: Tree[float]
        关键点坐标 XZY 次序      
    ConnectionStart: Tree[float]
        连接线起始点
    ConnectionEnd: Tree[float]
        连接线结束点 
    RkneeRangleLength: float
        实际提取的右小腿长度
    Landmarks: List[Point3D]
        关键点坐标
    ConnectionLines: List[Line]
        关键点连接线   
"""

import mediapipe as mp
import cv2
import math
from mediapipe.tasks.python import vision

import ghpythonlib.treehelpers as th
import rhinoscriptsyntax as rs

ghenv.Component.Name = "MediaPipePoseImage"
ghenv.Component.NickName = "MediaPipePoseImage"
ghenv.Component.Description = (
    "用 MediaPipe 库的人体姿态估计提取人体关键点信息，包括关键点之间的连线"
)
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 1


def distance_3d(point1, point2):
    """
    计算三维空间中两点之间的欧氏距离。

    Args:
        point1 (tuple): 第一个点为 (x1, y1, z1)
        point2 (tuple): 第二个点为 (x2, y2, z2)

    Returns:
        float: 两点之间的距离
    """
    x1, y1, z1 = point1
    x2, y2, z2 = point2
    return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2 + (z2 - z1) ** 2)


def poseDetection(
    static_image_mode=True, model_complexity=2, min_detection_confidence=0.3
):
    # 初始化 MediaPipe 姿势（Pose）
    mp_pose = mp.solutions.pose
    image = cv2.imread(ImgPath)

    # 2. 运行 MediaPipe 的 Pose
    with mp_pose.Pose(
        static_image_mode=True, min_detection_confidence=0.3, model_complexity=2
    ) as pose:
        results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        if results.pose_world_landmarks:
            nose_world_landmark = results.pose_world_landmarks.landmark[
                mp_pose.PoseLandmark.NOSE
            ]
            wlandmarks = results.pose_world_landmarks.landmark

            x = [wlandmark.x for wlandmark in wlandmarks]
            y = [wlandmark.y for wlandmark in wlandmarks]
            z = [wlandmark.z for wlandmark in wlandmarks]
            xyz = list(zip(x, y, z))

            XYZ = th.list_to_tree(xyz)
            RkneeRangleLength = distance_3d(xyz[26], xyz[28])
            scale = kneeAngleLength / RkneeRangleLength

            X = [i * scale for i in x]
            Y = [i * scale for i in z]
            Z = [-i * scale + HeightRef for i in y]
            xyz = list(zip(X, Y, Z))
            XZY = th.list_to_tree(xyz)

            connection_start = []
            connection_end = []
            ConnectionLines = []
            for connection in vision.PoseLandmarksConnections.POSE_LANDMARKS:
                cx_s = wlandmarks[connection.start].x
                cy_s = wlandmarks[connection.start].y
                cz_s = wlandmarks[connection.start].z
                pt_s = [cx_s * scale, cz_s * scale, -cy_s * scale + HeightRef]
                connection_start.append(pt_s)

                cx_e = wlandmarks[connection.end].x
                cy_e = wlandmarks[connection.end].y
                cz_e = wlandmarks[connection.end].z
                pt_e = [cx_e * scale, cz_e * scale, -cy_e * scale + HeightRef]
                connection_end.append(pt_e)

                ConnectionStart = th.list_to_tree(connection_start)
                ConnectionEnd = th.list_to_tree(connection_end)

                ConnectionLines.append(rs.AddLine(pt_s, pt_e))

            Landmarks = rs.AddPoints(xyz)

        else:
            print("No pose world landmarks detected.")

    return (
        X,
        Y,
        Z,
        XYZ,
        XZY,
        ConnectionStart,
        ConnectionEnd,
        RkneeRangleLength,
        Landmarks,
        ConnectionLines,
    )


if __name__ == "__main__":
    try:
        (
            X,
            Y,
            Z,
            XYZ,
            XZY,
            ConnectionStart,
            ConnectionEnd,
            RkneeRangleLength,
            Landmarks,
            ConnectionLines,
        ) = poseDetection(StaticImageMode, ModelComplexity, MinDetectionConfidence)
    except Exception as e:
        print(e)
  

  """
用 MediaPipe 库的人体姿态估计提取视频中人体关键点信息，包括关键点之间的连线。
Inputs:
    Videopath: Item Access; str
        视频文件路径
    MinDetectionConfidence: Item Access; float
        人体检测模型（person-detection model）最小置信度，值域为 [0.0,1.0]
    MinTrackingConfidence: Item Access; float
        关键点追踪模型（landmark-tracking model ）最小置信度，值域为 [0.0,1.0]
    FrameInterval: Item Access; int
        间隔帧数
    HeightRef: Item Access; float
        参考位置高度
    kneeAngleLength: Item Access; float
        由小腿长度，用于控制人体身高
Output:
    LandmarksAnim: Tree[Point3D]
        视频各帧图像关键点坐标
    ConnectionLinesAnime: Tree[Line]
        视频各帧图像关键点连接线    

"""

import cv2
import mediapipe as mp
import math
from mediapipe.tasks.python import vision

import ghpythonlib.treehelpers as th
import rhinoscriptsyntax as rs

ghenv.Component.Name = "MediaPipePoseVideo"
ghenv.Component.NickName = "MediaPipePoseVideo"
ghenv.Component.Description = "用 MediaPipe 库的人体姿态估计提取视频中人体关键点信息，包括关键点之间的连线"
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 1


def distance_3d(point1, point2):
    """同 MediaPipePoseImage 组件中 distance_3d() 定义"""

def poseDetection_video(
    min_detection_confidence=0.5, min_tracking_confidence=0.5
):
    # 打开视频文件
    cap = cv2.VideoCapture(Videopath)
    # 初始化 MediaPipe 姿势（Pose）
    mp_pose = mp.solutions.pose

    # 运行 MediaPipe 的 Pose
    with mp_pose.Pose(
        static_image_mode=False,
        min_detection_confidence=min_detection_confidence,
        min_tracking_confidence=min_tracking_confidence,
    ) as pose:
        frame_count = 0
        LandmarksAnim = []
        ConnectionLinesAnime = []
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break  # 如果没有更多的帧可用，则中断循环

            # 只处理基于输入参数 FrameInterval 的第 n 帧
            if frame_count % FrameInterval == 0:
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb.flags.writeable = False  
                # 处理帧图像并提取人体姿态估计关键点坐标及其连接线
                results = pose.process(image_rgb)

                if results.pose_world_landmarks:
                    nose_world_landmark = results.pose_world_landmarks.landmark[
                        mp_pose.PoseLandmark.NOSE
                    ]

                    wlandmarks = results.pose_world_landmarks.landmark
                    x = [wlandmark.x for wlandmark in wlandmarks]
                    y = [wlandmark.y for wlandmark in wlandmarks]
                    z = [wlandmark.z for wlandmark in wlandmarks]
                    xyz = list(zip(x, y, z))

                    RkneeRangleLength = distance_3d(xyz[26], xyz[28])
                    scale = kneeAngleLength / RkneeRangleLength

                    X = [i * scale for i in x]
                    Y = [i * scale for i in z]
                    Z = [-i * scale + HeightRef for i in y]
                    xyz = list(zip(X, Y, Z))
                    XYZ = th.list_to_tree(xyz)

                    connection_start = []
                    connection_end = []
                    ConnectionLines = []
                    for (
                        connection
                    ) in vision.PoseLandmarksConnections.POSE_LANDMARKS:
                        cx_s = wlandmarks[connection.start].x
                        cy_s = wlandmarks[connection.start].y
                        cz_s = wlandmarks[connection.start].z
                        pt_s = [
                            cx_s * scale,
                            cz_s * scale,
                            -cy_s * scale + HeightRef,
                        ]
                        connection_start.append(pt_s)

                        cx_e = wlandmarks[connection.end].x
                        cy_e = wlandmarks[connection.end].y
                        cz_e = wlandmarks[connection.end].z
                        pt_e = [
                            cx_e * scale,
                            cz_e * scale,
                            -cy_e * scale + HeightRef,
                        ]
                        connection_end.append(pt_e)

                        ConnectionStart = th.list_to_tree(connection_start)
                        ConnectionEnd = th.list_to_tree(connection_end)

                        ConnectionLines.append(rs.AddLine(pt_s, pt_e))

                    Landmarks = rs.AddPoints(xyz)
                    LandmarksAnim.append(Landmarks)
                    ConnectionLinesAnime.append(ConnectionLines)

            frame_count += 1

    LandmarksAnim = th.list_to_tree(LandmarksAnim)
    ConnectionLinesAnime = th.list_to_tree(ConnectionLinesAnime)

    return LandmarksAnim, ConnectionLinesAnime


if __name__ == "__main__":
    try:
        LandmarksAnim, ConnectionLinesAnime = poseDetection_video(
            MinDetectionConfidence, MinTrackingConfidence
        )
    except Exception as e:
        print(e)
  

  """
还原视频中用 MediaPipe 库的人体姿态估计提取视频中人体关键点信息和连接线的动画

Inputs:
    LandmarksAnim: Tree[Point3D]
        连续帧图像关键点
    ConnectionLinesAnime: Tree[Line]
        连续帧图像关键点连接线
    FramesStart: Item[int]
        起始帧
    FramesEnd: Item[int]
        结束帧
    FramesInterval: Item[int]
        帧间隔数
    Direction: Item[Vector3D]
        移动方向
    Stride：Item[float]
        每步移动的距离
    Stride: Item[float]
        更新速度
    Reset: Item[bool]
        重启计时
    Running: Item[bool]
        执行动画
"""

import Grasshopper as gh
import ghpythonlib.treehelpers as th
import numpy as np
import rhinoscriptsyntax as rs

ghenv.Component.Name = "MediaPipePoseAnim"
ghenv.Component.NickName = "MediaPipePoseAnim"
ghenv.Component.Description = "还原视频中用 MediaPipe 库的人体姿态估计提取视频中人体关键点信息和连接线的动画"
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 1


def move_points_along_vector_3d(points, vector_2d, distance):
    """
    在 XY 平面上沿2D向量连续移动3D点

    Args:
        points (list of tuples):初始3D点的列表，点格式为 (x, y, z).
        vector_2d (tuple): 一个2D向量(vx, vy)，定义了 XY 平面上的移动方向
        distance (float): 矢量方向每一步移动的距离
        steps (int): 移动点的步数

    Returns:
        list of list: 包含每步更新点的列表。
    """
    # 二维矢量归一化，得到单位方向矢量
    vx, vy = vector_2d
    magnitude = np.sqrt(vx**2 + vy**2)
    direction = (vx / magnitude, vy / magnitude)

    # 用于存储点随时间变化的新位置
    trajectory = []

    # 沿给定方向，按每步移动的距离移动点
    for i in range(len(points)):
        # 更新列表中的每个点
        new_positions = [
            [
                float(rs.PointCoordinates(pt)[0] + direction[0] * distance * i),
                float(rs.PointCoordinates(pt)[1] + direction[1] * distance * i),
                rs.PointCoordinates(pt)[2],
            ]
            for pt in points[i]
        ]
        trajectory.append(rs.AddPoints(new_positions))

    return trajectory


def updateComponent(Speed):
    """更新此组件类似于使用 GH 定时器"""

    def callBack(e):
        ghenv.Component.ExpireSolution(False)

    ghenv.Component.OnPingDocument().ScheduleSolution(
        Speed, gh.Kernel.GH_Document.GH_ScheduleDelegate(callBack)
    )


LandmarksAnim = th.tree_to_list(LandmarksAnim)
ConnectionLinesAnime = th.tree_to_list(ConnectionLinesAnime)

print("+++",Direction)

if Direction is not None:
    vector_2d = [Direction[0], Direction[1]]
    LandmarksAnim = move_points_along_vector_3d(
        LandmarksAnim, vector_2d, Stride
    )
    start_pts = [
        [rs.AddPoint(rs.CurveStartPoint(p)) for p in g]
        for g in ConnectionLinesAnime
    ]
    start_pts_moved = move_points_along_vector_3d(start_pts, vector_2d, Stride)
    end_pts = [
        [rs.AddPoint(rs.CurveEndPoint(p)) for p in g]
        for g in ConnectionLinesAnime
    ]
    end_pts_moved = move_points_along_vector_3d(end_pts, vector_2d, Stride)
    ConnectionLinesAnime = [
        [rs.AddLine(s, e) for s, e in zip(sg, eg)]
        for sg, eg in zip(start_pts_moved, end_pts_moved)
    ]

frames = list(range(FramesStart, FramesEnd, FramesInterval))

# 实例化或重置计数器变量
if "count" not in globals() or Reset:
    count = FramesStart

# 更新变量和组件
if Running and not Reset and count < FramesEnd:
    count += FramesInterval
    updateComponent(Speed)
    Landmarks = LandmarksAnim[count]
    ConnectionLines = ConnectionLinesAnime[count]

    if count >= FramesEnd:
        count = FramesStart

# 输出计数
Counter = count
  

  """
读取 AGORA 人体姿态估计数据集，使用 SMPL-X （Skinned Multi-Person Linear Model ）蒙皮多人线性模型重建三维人体模型
Inputs:
    SMPLXModelPath: Item[str]
        SMPL-X 蒙皮多人线性模型文件路径
    AgoraBodyPath:Item[str]
         AGORA 人体姿态估计数据集文件路径
    Idx: Item[int]
        数据集数据索引
    ModelType: Item[str]
        为 SMPL 系列模型，包括 smpl、smplx、smplh 等
    Gender: Item[str]
        加载哪个性别的模型，包括 neutral、male 和female 等
Output:
    DatasetKeys: List[str]
        AGORA 人体姿态估计数据集数据名
    SampleLength: int
        样本数
    Imgname: str
        对应输入端 Idx 参数的数据集样本图像名
    Vertices: List[Point3D]
        蒙骗控制点
    Faces: Tree[int]
        蒙骗控制点的单元面索引（三边面）
    Joints: List[Point3D]
        人体姿态估计关键点
"""

import numpy as np
import torch
from smplx import SMPLX
import rhinoscriptsyntax as rs
import ghpythonlib.treehelpers as th

ghenv.Component.Name = "AGORA-SMPLXHumanPose"
ghenv.Component.NickName = "AGORA-SMPLXHumanPose"
ghenv.Component.Description = "读取 AGORA 人体姿态估计数据集，使用 SMPL-X （Skinned Multi-Person Linear Model ）蒙皮多人线性模型重建三维人体模型"
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 1


def load_agora_parameters(npz_file_path):
    """
    Load SMPL-X parameters from an AGORA .npz file.

    Args:
        npz_file_path (str): Path to the agora-body.npz file.

    Returns:
        dict: Dictionary containing SMPL-X parameters like 'betas', 'body_pose', 'global_orient', and 'transl'.
    """
    # 加载 AGORA 数据集 .npz 文件
    data = np.load(npz_file_path)

    # 从 .npz 文件中提取必要的 SMPL-X 参数
    smplx_params = {
        "betas": torch.tensor(
            data["shape"], dtype=torch.float32
        ),  # 形状参数 Shape parameters (10D or more)
        "body_pose": torch.tensor(
            data["pose_world"], dtype=torch.float32
        ),  # 姿态参数 Body pose parameters (21*3 = 63D)
        "global_orient": torch.tensor(
            data["trans_world"], dtype=torch.float32
        ),  # 全局朝向 (3D)
        "transl": torch.tensor(
            data["trans_world"], dtype=torch.float32
        ),  # 在 3D 空间中的变换 (3D)
    }
    return smplx_params


# 创建自定义人体姿势的函数
def create_custom_pose(
    smpl_model_path,
    body_pose_params,
    global_orient_params,
    betas_params=None,
    transl=None,
    model_type="smplx",
    gender="neutral",
    batch_size=1,
):
    # 加载 SMPL-X 模型
    smplx_model = SMPLX(
        smpl_model_path, model_type=model_type, gender=gender, batch_size=batch_size
    )

    # 将参数转换为张量
    body_pose = torch.tensor(body_pose_params, dtype=torch.float32).unsqueeze(0)
    global_orient = torch.tensor(
        global_orient_params, dtype=torch.float32
    ).unsqueeze(0)
    transl = torch.tensor(transl,dtype=torch.float32).unsqueeze(0)

    # 如果没有提供形状参数，则使用零姿态（模板网格）
    if betas_params is None:
        betas = torch.zeros([1, 10], dtype=torch.float32)
    else:
        betas = torch.tensor(betas_params, dtype=torch.float32).unsqueeze(0)

    # 使用提取的 AGORA 姿态参数生成SMPL-X模型输出
    output = smplx_model(
        body_pose=body_pose, global_orient=global_orient, betas=betas, transl=transl
    )

    # 提取蒙皮顶点和关键点
    vertices = output.vertices.detach().cpu().numpy().squeeze()
    joints = output.joints.detach().cpu().numpy().squeeze()

    faces = smplx_model.faces

    return vertices, joints, faces


if __name__ == "__main__":
    if ModelType is None:
        ModelType = "smplx"
    if Gender is None:
        Gender = "neutral"

    try:
        data = np.load(AgoraBodyPath)
        DatasetKeys = data.files

        smplx_parameters = load_agora_parameters(AgoraBodyPath)
        custom_body_pose = smplx_parameters["body_pose"][Idx].reshape(55, 3)[
            1:22
        ]
        custom_global_orient = smplx_parameters["global_orient"][Idx]
        betasParams = smplx_parameters["betas"][Idx][:10]
        transl = smplx_parameters["transl"][Idx]

        imgnames = data["imgname"]
        Imgname = f"{imgnames[Idx]}"
        SampleLength = len(imgnames)

        vertices, joints, faces = create_custom_pose(
            SMPLXModelPath,
            custom_body_pose,
            custom_global_orient,
            betasParams,
            model_type=ModelType,
            gender=Gender,
            transl=transl
        )
        XYZ_Trans = [0, 2, 1]
        vertices = vertices[:, XYZ_Trans]
        Vertices = rs.AddPoints(vertices)
        Faces = th.list_to_tree(faces.tolist())
        joints = joints[:, XYZ_Trans]
        Joints = rs.AddPoints(joints)
    except Exception as e:
        print(e)
  

  """
使用 VPoser （Variational Human Pose Prior ）先验分布模型提取 SMPL-X 的姿态参数（pose）后建立 3D 人体模型。除了在 Python Script 环境下安装 SMPL-X 库，也需要按照 omegaconf 和 loguru库
Inputs:
    VposerModel: Item[str]
        Variational Human Pose Prior（VPoser）模型文件路径
    SMPLXModel: Item[str]
        SMPL-X 模型文件路径
    NumPoses: Item[int]
        生成连续姿态人体模型的数量
    BetasParams: Iist[float]
        形状成分参数
    Seed: Item[int]
        随机种子
Output:
    Vertices: Tree|List[Point3D]
        人体模型蒙皮顶点。如果 NumPoses 参数不为0，则输出为 Tree；否则为 List
    Faces: Tree[int]
        人体模型蒙皮顶点三边面索引
    Joints: List[Point3D]
        人体模型关键点列表
"""

import numpy as np
import torch
from human_body_prior.tools.model_loader import load_model
from human_body_prior.models.vposer_model import VPoser
from smplx import SMPLX
import rhinoscriptsyntax as rs
import ghpythonlib.treehelpers as th

print(f"cuda.is_available:{torch.cuda.is_available()}")

ghenv.Component.Name = "VPoserSMPL-X"
ghenv.Component.NickName = "VPoserSMPL-X"
ghenv.Component.Description = "使用 VPoser （Variational Human Pose Prior ）先验分布模型提取 SMPL-X 的姿态参数（pose）后建立 3D 人体模型"
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 1

def create_custom_pose(
    smpl_model_path,
    body_pose_params,
    global_orient_params,
    betas_params=None,
    model_type="smplx",
    gender="neutral",
    batch_size=1,
):
    '''同 AGORA-SMPLXHumanPose 组件定义的 create_custom_pose 函数'''
    return vertices, joints, faces


def vposeModel2SMPLx():
    vp, ps = load_model(
        VposerModel,
        model_code=VPoser,
        remove_words_in_model_weights="vp_model.",
        disable_grad=True,
    )
    vp = vp.to("cuda")
    custom_global_orient = np.zeros(3)
    if BetasParams is None:
        betasParams = np.zeros(10)
    else:
        betasParams = np.array(BetasParams)
    XYZ_transpose = [0, 2, 1]

    print("++#", NumPoses)

    if NumPoses is None or NumPoses < 1:
        # 从正态分布中采样一个 32 维向量
        poZ_body_sample = torch.from_numpy(
            np.random.randn(1, 32).astype(np.float32)
        ).to("cuda")
        pose_body = (
            vp.decode(poZ_body_sample)["pose_body"].contiguous().view(-1, 63)
        )

        print(
            f"poZ_body_sample.shape={poZ_body_sample.shape}"
        )  # torch.Size([1, 32])
        print(f"pose_body.shape={pose_body.shape}")  # torch.Size([1, 63])

        custom_body_pose = pose_body[0].reshape(21, 3).cpu().numpy()
        vertices, joints, faces = create_custom_pose(
            SMPLXModel,
            custom_body_pose,
            custom_global_orient,
            betasParams,
        )
        vertices = vertices[:, XYZ_transpose]
        Vertices = rs.AddPoints(vertices)
        Faces = th.list_to_tree(faces.tolist())

        joints = joints[:, XYZ_transpose]
        Joints = rs.AddPoints(joints)
    else:
        sampled_pose_body = (
            vp.sample_poses(num_poses=NumPoses)["pose_body"]
            .contiguous()
            .view(NumPoses, -1)
        )  # 将生成一个 Nx1x21x3 维度的人体姿态向量
        Vertices, Faces, Joints = [], [], []
        for sample_pose in sampled_pose_body:
            vertices, joints, faces = create_custom_pose(
                SMPLXModel,
                sample_pose.cpu().numpy(),
                custom_global_orient,
                betasParams,
            )
            vertices = vertices[:, XYZ_transpose]
            Vertices_g = rs.AddPoints(vertices)

            joints = joints[:, XYZ_transpose]
            Joints_g = rs.AddPoints(joints)

            Vertices.append(Vertices_g)
            Faces.append(faces.tolist())
            Joints.append(Joints_g)

        Vertices = th.list_to_tree(Vertices)
        Faces = th.list_to_tree(Faces)
        Joints = th.list_to_tree(Joints)

    return Vertices, Faces, Joints


if __name__ == "__main__":
    try:
        Vertices, Faces, Joints = vposeModel2SMPLx()
    except Exception as e:
        print(e)
  

  '''
对单幅 RGB 图像，相机空间中的多人进行人体 3D 建模（使用 SMPL 模型）。迁移代码地址：https://github.com/naver/multi-hmr
Inputs:
    SMPLX_DIR: Item[str]
        SMPLX 模型所在文件夹路径名
    MEAN_PARAMS: Item[str]
        smpl_mean_params.npz 文件路径。下载地址为：https://openmmlab-share.oss-cn-hangzhou.aliyuncs.com/mmhuman3d/models/smpl_mean_params.npz?versionId=CAEQHhiBgICN6M3V6xciIDU1MzUzNjZjZGNiOTQ3OWJiZTJmNThiZmY4NmMxMTM4
    MultiHMR: Item[str]
        Multi-HMR 预训练模型，下载地址：https://github.com/naver/multi-hmr
    ImgPath: Item[str]
        图像文件路径名
Output:
    humans: List[Dict]
        图像中所有识别提取的人体模型参数，包含 SMPLX 人体姿态估计的参数
    l_duration: List[floatt]
        模型计算所花费的时间
'''

import os

os.environ["PYOPENGL_PLATFORM"] = "egl"
os.environ["EGL_DEVICE_ID"] = "0"

from argparse import ArgumentParser
import random
import numpy as np
from PIL import Image, ImageOps
import torch
import time

from moths.multi_hmr.utils import (
    normalize_rgb,
    get_focalLength_from_fieldOfView,
)
from moths.multi_hmr import Model

ghenv.Component.Name = "Multi-HMR"
ghenv.Component.NickName = "Multi-HMR"
ghenv.Component.Description = "对单幅 RGB 图像，相机空间中的多人进行人体 3D 建模"
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 1

torch.cuda.empty_cache()
np.random.seed(seed=0)
random.seed(0)


def open_image(img_path, img_size, device=torch.device("cuda")):
    """打开图像，调整大小并填充"""
    img_pil = Image.open(img_path).convert("RGB")
    img_pil = ImageOps.contain(
        img_pil, (img_size, img_size)
    )  # 保持相同的纵横比

    # 为可视化保留一份副本
    img_pil_bis = ImageOps.pad(
        img_pil.copy(), size=(img_size, img_size), color=(255, 255, 255)
    )
    img_pil = ImageOps.pad(
        img_pil, size=(img_size, img_size)
    )  # 用值 0 填充

    resize_img = np.asarray(img_pil)

    # 标准化
    resize_img = normalize_rgb(resize_img)
    x = torch.from_numpy(resize_img).unsqueeze(0).to(device)
    return x, img_pil_bis


def get_camera_parameters(
    img_size, fov=60, p_x=None, p_y=None, device=torch.device("cuda")
):
    """给定图像大小、视场（fov）和主点坐标（principal point coordinates），返回相机参数矩阵 K"""
    K = torch.eye(3)
    # 获取焦距（focal length）
    focal = get_focalLength_from_fieldOfView(fov=fov, img_size=img_size)
    K[0, 0], K[1, 1] = focal, focal

    # 设置主点（principal point）
    if p_x is not None and p_y is not None:
        K[0, -1], K[1, -1] = p_x * img_size, p_y * img_size
    else:
        K[0, -1], K[1, -1] = img_size // 2, img_size // 2

    # 增加维度
    K = K.unsqueeze(0).to(device)
    return K


def load_model(
    model_name, SMPLX_DIR, MEAN_PARAMS, multiHMR, device=torch.device("cuda")
):
    """打开模型快照（checkpoint），使用保存的参数构建 Multi-HMR，加载模型权重"""
    ckpt = torch.load(multiHMR, map_location=device)

    # 获取保存在模型快照中的参数来重建模型
    kwargs = {}
    for k, v in vars(ckpt["args"]).items():
        kwargs[k] = v

    # 建立模型.
    kwargs["type"] = ckpt["args"].train_return_type
    kwargs["img_size"] = ckpt["args"].img_size[0]
    kwargs["SMPLX_DIR"] = SMPLX_DIR
    kwargs["MEAN_PARAMS"] = MEAN_PARAMS
    model = Model(**kwargs).to(device)

    # 向模型中加载权重
    model.load_state_dict(ckpt["model_state_dict"], strict=False)
    print("Weights have been loaded")

    return model


def forward_model(
    model,
    input_image,
    camera_parameters,
    det_thresh=0.3,
    nms_kernel_size=1,
):
    """对输入的图像和相机参数进行前向传播"""

    # 模型的前向传播.
    with torch.no_grad():
        with torch.cuda.amp.autocast(enabled=True):
            humans = model(
                input_image,
                is_training=False,
                nms_kernel_size=int(nms_kernel_size),
                det_thresh=det_thresh,
                K=camera_parameters,
            )

    return humans

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--model_name", type=str, default="multiHMR_672_S"
    )  # 'multiHMR_896_L_synth'
    parser.add_argument("--img_folder", type=str, default="example_data")
    parser.add_argument("--out_folder", type=str, default="demo_out")
    parser.add_argument("--save_mesh", type=int, default=0, choices=[0, 1])
    parser.add_argument("--extra_views", type=int, default=0, choices=[0, 1])
    parser.add_argument("--det_thresh", type=float, default=0.3)
    parser.add_argument("--nms_kernel_size", type=float, default=3)
    parser.add_argument("--fov", type=float, default=60)
    parser.add_argument(
        "--distance",
        type=int,
        default=0,
        choices=[0, 1],
        help="add distance on the reprojected mesh",
    )
    parser.add_argument(
        "--unique_color",
        type=int,
        default=0,
        choices=[0, 1],
        help="only one color for all humans",
    )

    args = parser.parse_args()
    dict_args = vars(args)
    assert torch.cuda.is_available()

    try:
        model = load_model(args.model_name, SMPLX_DIR, MEAN_PARAMS, MultiHMR)
        l_duration = []
        img_size = model.img_size
        x, img_pil_nopad = open_image(
            os.path.join(args.img_folder, ImgPath), img_size
        )

        # Get camera parameters
        p_x, p_y = None, None
        K = get_camera_parameters(model.img_size, fov=args.fov, p_x=p_x, p_y=p_y)

        # Make model predictions
        start = time.time()
        humans = forward_model(
            model,
            x,
            K,
            det_thresh=args.det_thresh,
            nms_kernel_size=args.nms_kernel_size,
        )
        duration = time.time() - start
        l_duration.append(duration)
    except Exception as e:
        print(e)
  

  """
提取 Multi-HMR 组件输出端 humans 中，对应 SMPLX 的参数，并建立 3D 人体模型
Inputs:
    Humans: List[Dict]
        Multi-HMR 组件输出端 humans 数据
    SMPLXModelPath: Item[str]
        SMPLX 模型文件路径
    XYZ_trans: Item[str]
        X,Y,Z 坐标置换，例如 [0,2,1]
    Idx: Item[int]
        提取对象索引值
Output:
    Vertices: List[Point3D]
        蒙骗控制点
    Faces: Tree[int]
        蒙骗控制点的单元面索引（三边面）
    Joints: List[Point3D]
        人体姿态估计关键点    
"""

from smplx import SMPLX
import rhinoscriptsyntax as rs
import ghpythonlib.treehelpers as th
import torch

ghenv.Component.Name = "Multi-HMR2SMPLX"
ghenv.Component.NickName = "Multi-HMR2SMPLX"
ghenv.Component.Description = "提取 Multi-HMR 组件输出端 humans 中，对应 SMPLX 的参数，并建立 3D 人体模型"
ghenv.Component.Message = "0.0.1"
ghenv.Component.Category = "Moths"
ghenv.Component.SubCategory = "Human Pose Estimation"
ghenv.Component.PanelSection = 2

def create_custom_pose(
    smpl_model_path,
    body_pose_params,
    global_orient_params,
    betas_params=None,
    transl=None,
    model_type="smplx",
    gender="neutral",
    batch_size=1,
):
    '''同前文定义的 create_custom_pose() 函数'''
    return vertices, joints, faces


if __name__ == "__main__":
    smplxParams = Humans[Idx]
    custom_body_pose = smplxParams["rotvec"].cpu().numpy()[1:22]
    custom_global_orient = smplxParams["transl_pelvis"].cpu().numpy()
    betasParams = smplxParams["shape"].cpu().numpy()

    vertices, joints, faces = create_custom_pose(
        SMPLXModelPath, custom_body_pose, custom_global_orient, betasParams
    )

    if XYZ_trans is not None:
        orientation = eval(XYZ_trans)
    else:
        orientation = [0, 1, 2]

    vertices = vertices[:, orientation]
    vertices[:, 2] = -vertices[:, 2]

    Vertices = rs.AddPoints(vertices)
    Faces = th.list_to_tree(faces.tolist())

    joints = joints[:, orientation]
    joints[:, 2] = -joints[:, 2]
    Joints = rs.AddPoints(joints)