mGPT/render/blender/joints.py

import math

import bpy
import numpy as np

from mGPT.utils.joints import (humanml3d_joints, humanml3d_kinematic_tree,
                               mmm_joints, mmm_kinematic_tree,
                               mmm_to_smplh_scaling_factor)

# from .materials import colored_material_diffuse_BSDF as colored_material
from .materials import colored_material_relection_BSDF as colored_material

sat_factor = 1.1

JOINTS_MATS = [
    # colored_material(0.2500, 0.0357, 0.0349, saturation_factor = sat_factor),
    # # colored_material(0.4500, 0.0357, 0.0349),
    # colored_material(0.6500, 0.175, 0.0043, saturation_factor = sat_factor),
    # colored_material(0.0349, 0.3500, 0.0349, saturation_factor = sat_factor),
    # colored_material(0.018, 0.059, 0.600, saturation_factor = sat_factor),
    # colored_material(0.032, 0.325, 0.421, saturation_factor = sat_factor),
    # colored_material(0.3, 0.3, 0.3, saturation_factor = sat_factor),
    colored_material(0.3500, 0.0357, 0.0349, saturation_factor=sat_factor),
    # colored_material(0.4500, 0.0357, 0.0349),
    colored_material(0.6500, 0.175, 0.0043, saturation_factor=sat_factor),
    colored_material(0.0349, 0.3500, 0.0349, saturation_factor=sat_factor),
    colored_material(0.018, 0.059, 0.600, saturation_factor=sat_factor),
    colored_material(0.032, 0.325, 0.421, saturation_factor=sat_factor),
    colored_material(0.3, 0.3, 0.3, saturation_factor=sat_factor),
]


class Joints:

    def __init__(self,
                 data,
                 *,
                 mode,
                 canonicalize,
                 always_on_floor,
                 jointstype="mmm",
                 **kwargs):
        data = prepare_joints(
            data,
            canonicalize=canonicalize,
            always_on_floor=always_on_floor,
            jointstype=jointstype,
        )

        self.data = data
        self.mode = mode

        self.N = len(data)

        self.N = len(data)
        self.trajectory = data[:, 0, [0, 1]]

        if jointstype == "mmm":
            self.kinematic_tree = mmm_kinematic_tree
            self.joints = mmm_joints
            self.joinst.append("")
        elif jointstype == "humanml3d":
            self.kinematic_tree = humanml3d_kinematic_tree
            self.joints = humanml3d_joints

        self.mat = JOINTS_MATS

    def get_sequence_mat(self, frac):
        return self.mat

    def get_root(self, index):
        return self.data[index][0]

    def get_mean_root(self):
        return self.data[:, 0].mean(0)

    def load_in_blender(self, index, mats):
        skeleton = self.data[index]
        head_mat = mats[0]
        body_mat = mats[-1]
        for lst, mat in zip(self.kinematic_tree, mats):
            for j1, j2 in zip(lst[:-1], lst[1:]):
                # spine and head
                if self.joints[j2] in [
                        "BUN",
                ]:
                    sphere_between(skeleton[j1], skeleton[j2], head_mat)
                elif self.joints[j2] in [
                        "LE",
                        "RE",
                        "LW",
                        "RW",
                ]:
                    cylinder_sphere_between(skeleton[j1], skeleton[j2], 0.040,
                                            mat)
                elif self.joints[j2] in [
                        "LMrot",
                        "RMrot",
                        "RK",
                        "LK",
                ]:
                    cylinder_sphere_between(skeleton[j1], skeleton[j2], 0.040,
                                            mat)
                elif self.joints[j2] in [
                        "LS",
                        "RS",
                        "LF",
                        "RF",
                ]:
                    cylinder_between(skeleton[j1], skeleton[j2], 0.040, mat)
                elif self.joints[j2] in ["RK", "LK"]:
                    print(self.joints[j1], self.joints[j2])
        # body
        sphere(0.14, skeleton[self.joints.index("BLN")], body_mat)
        sphere_between(
            skeleton[self.joints.index("BLN")],
            skeleton[self.joints.index("root")],
            body_mat,
            factor=0.28,
        )
        sphere(0.11, skeleton[self.joints.index("root")], body_mat)
        # sphere_between(
        #     skeleton[self.joints.index("BLN")],
        #     skeleton[self.joints.index("BT")],
        #     mats[0],
        # )
        # hip
        # sphere_between(
        #     skeleton[self.joints.index("LH")],
        #     skeleton[self.joints.index("RH")],
        #     mats[0],
        #     factor=0.6,
        # )
        #
        # sphere(skeleton[self.joints.index("BLN")], 0.05, mats[0])
        # sphere_between(skeleton[13], skeleton[14], mat)
        # node
        # print(self.joints.index("BUN"))
        # print(len(lst))
        # sphere(lst[self.joints.index("BUN")], 0.2, mat)  # head

        return ["Cylinder", "Sphere"]

    def __len__(self):
        return self.N


def softmax(x, softness=1.0, dim=None):
    maxi, mini = x.max(dim), x.min(dim)
    return maxi + np.log(softness + np.exp(mini - maxi))


def softmin(x, softness=1.0, dim=0):
    return -softmax(-x, softness=softness, dim=dim)


def get_forward_direction(poses, jointstype="mmm"):
    if jointstype == "mmm" or jointstype == "mmmns":
        joints = mmm_joints
    elif jointstype == "humanml3d":
        joints = humanml3d_joints
    else:
        raise TypeError("Only supports mmm, mmmns and humanl3d jointstype")
    # Shoulders
    LS, RS = joints.index("LS"), joints.index("RS")
    # Hips
    LH, RH = mmm_joints.index("LH"), mmm_joints.index("RH")

    across = (poses[..., RH, :] - poses[..., LH, :] + poses[..., RS, :] -
              poses[..., LS, :])
    forward = np.stack((-across[..., 2], across[..., 0]), axis=-1)
    forward = forward / np.linalg.norm(forward, axis=-1)
    return forward


def cylinder_between(t1, t2, r, mat):
    x1, y1, z1 = t1
    x2, y2, z2 = t2

    dx = x2 - x1
    dy = y2 - y1
    dz = z2 - z1
    dist = math.sqrt(dx**2 + dy**2 + dz**2)

    bpy.ops.mesh.primitive_cylinder_add(radius=r,
                                        depth=dist,
                                        location=(dx / 2 + x1, dy / 2 + y1,
                                                  dz / 2 + z1))

    phi = math.atan2(dy, dx)
    theta = math.acos(dz / dist)
    bpy.context.object.rotation_euler[1] = theta
    bpy.context.object.rotation_euler[2] = phi
    # bpy.context.object.shade_smooth()
    bpy.context.object.active_material = mat

    bpy.ops.mesh.primitive_uv_sphere_add(radius=r, location=(x1, y1, z1))
    bpy.context.object.active_material = mat
    bpy.ops.mesh.primitive_uv_sphere_add(radius=r, location=(x2, y2, z2))
    bpy.context.object.active_material = mat


def cylinder_sphere_between(t1, t2, r, mat):
    x1, y1, z1 = t1
    x2, y2, z2 = t2
    dx = x2 - x1
    dy = y2 - y1
    dz = z2 - z1
    dist = math.sqrt(dx**2 + dy**2 + dz**2)
    phi = math.atan2(dy, dx)
    theta = math.acos(dz / dist)
    dist = dist - 0.2 * r
    # sphere node
    sphere(r * 0.9, t1, mat)
    sphere(r * 0.9, t2, mat)
    # leveled cylinder
    bpy.ops.mesh.primitive_cylinder_add(
        radius=r,
        depth=dist,
        location=(dx / 2 + x1, dy / 2 + y1, dz / 2 + z1),
        enter_editmode=True,
    )
    bpy.ops.mesh.select_mode(type="EDGE")
    bpy.ops.mesh.select_all(action="DESELECT")
    bpy.ops.mesh.select_face_by_sides(number=32, extend=False)
    bpy.ops.mesh.bevel(offset=r, segments=8)
    bpy.ops.object.editmode_toggle(False)
    # bpy.ops.object.shade_smooth()
    bpy.context.object.rotation_euler[1] = theta
    bpy.context.object.rotation_euler[2] = phi
    bpy.context.object.active_material = mat


def sphere(r, t, mat):
    bpy.ops.mesh.primitive_uv_sphere_add(segments=50,
                                         ring_count=50,
                                         radius=r,
                                         location=t)
    # bpy.ops.mesh.primitive_uv_sphere_add(radius=r, location=t)
    # bpy.context.object.shade_smooth()
    bpy.context.object.active_material = mat


def sphere_between(t1, t2, mat, factor=1):
    x1, y1, z1 = t1
    x2, y2, z2 = t2

    dx = x2 - x1
    dy = y2 - y1
    dz = z2 - z1
    dist = math.sqrt(dx**2 + dy**2 + dz**2) * factor

    bpy.ops.mesh.primitive_uv_sphere_add(
        segments=50,
        ring_count=50,
        # bpy.ops.mesh.primitive_uv_sphere_add(
        radius=dist,
        location=(dx / 2 + x1, dy / 2 + y1, dz / 2 + z1))

    # bpy.context.object.shade_smooth()
    bpy.context.object.active_material = mat


def matrix_of_angles(cos, sin, inv=False):
    sin = -sin if inv else sin
    return np.stack((np.stack(
        (cos, -sin), axis=-1), np.stack((sin, cos), axis=-1)),
                    axis=-2)


def get_floor(poses, jointstype="mmm"):
    if jointstype == "mmm" or jointstype == "mmmns":
        joints = mmm_joints
    elif jointstype == "humanml3d":
        joints = humanml3d_joints
    else:
        raise TypeError("Only supports mmm, mmmns and humanl3d jointstype")
    # Feet
    LM, RM = joints.index("LMrot"), joints.index("RMrot")
    LF, RF = joints.index("LF"), joints.index("RF")
    ndim = len(poses.shape)

    foot_heights = poses[..., (LM, LF, RM, RF), 1].min(-1)
    floor_height = softmin(foot_heights, softness=0.5, dim=-1)
    return floor_height[tuple((ndim - 2) * [None])].T


def canonicalize_joints(joints, jointstype="mmm"):
    poses = joints.copy()

    translation = joints[..., 0, :].copy()

    # Let the root have the Y translation
    translation[..., 1] = 0
    # Trajectory => Translation without gravity axis (Y)
    trajectory = translation[..., [0, 2]]

    # Remove the floor
    poses[..., 1] -= get_floor(poses, jointstype)

    # Remove the trajectory of the joints
    poses[..., [0, 2]] -= trajectory[..., None, :]

    # Let the first pose be in the center
    trajectory = trajectory - trajectory[..., 0, :]

    # Compute the forward direction of the first frame
    forward = get_forward_direction(poses[..., 0, :, :], jointstype)

    # Construct the inverse rotation matrix
    sin, cos = forward[..., 0], forward[..., 1]
    rotations_inv = matrix_of_angles(cos, sin, inv=True)

    # Rotate the trajectory
    trajectory_rotated = np.einsum("...j,...jk->...k", trajectory,
                                   rotations_inv)

    # Rotate the poses
    poses_rotated = np.einsum("...lj,...jk->...lk", poses[..., [0, 2]],
                              rotations_inv)
    poses_rotated = np.stack(
        (poses_rotated[..., 0], poses[..., 1], poses_rotated[..., 1]), axis=-1)

    # Re-merge the pose and translation
    poses_rotated[..., (0, 2)] += trajectory_rotated[..., None, :]
    return poses_rotated


def prepare_joints(joints,
                   canonicalize=True,
                   always_on_floor=False,
                   jointstype="mmm"):
    # All face the same direction for the first frame
    if canonicalize:
        data = canonicalize_joints(joints, jointstype)
    else:
        data = joints

    # Rescaling, shift axis and swap left/right
    if jointstype == "humanml3d":
        data = data * mmm_to_smplh_scaling_factor
        data[..., 1] = - data[..., 1]

    # Swap axis (gravity=Z instead of Y)
    data = data[..., [2, 0, 1]]

    if jointstype == "mmm":
        # Make left/right correct
        data[..., [1]] = -data[..., [1]]

    # Center the first root to the first frame
    data -= data[[0], [0], :]

    # Remove the floor
    data[..., 2] -= data[..., 2].min()

    # Put all the body on the floor
    if always_on_floor:
        data[..., 2] -= data[..., 2].min(1)[:, None]

    return data


def NormalInDirection(normal, direction, limit=0.5):
    return direction.dot(normal) > limit


def GoingUp(normal, limit=0.5):
    return NormalInDirection(normal, (0, 0, 1), limit)


def GoingDown(normal, limit=0.5):
    return NormalInDirection(normal, (0, 0, -1), limit)


def GoingSide(normal, limit=0.5):
    return GoingUp(normal, limit) == False and GoingDown(normal,
                                                         limit) == False