inference.py

"""
推理逻辑：SMPL-X AMASS NPZ → G1 DOF/root_trans/root_rot

从 tools/inference.py 提取，路径适配 HF demo 目录结构。
"""
import os
import time
import torch
import numpy as np
from copy import deepcopy
from scipy import signal as sp_signal
from smplx import SMPLX

from mmengine.registry import MODELS

from src.utils.rotation_conversions import (
    axis_angle_to_matrix, matrix_to_axis_angle, matrix_to_rotation_6d,
    rotation_6d_to_matrix, matrix_to_quaternion,
)

# 确保 mmengine 注册表加载
import src  # noqa

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# ---------- 模型配置（内联，不依赖外部 .py 文件） ----------

MODEL_CFG = dict(
    init_cfg=None,
    n_embd=512,
    smplx_vqvae_cfg=dict(
        decoder_cfg=dict(
            activation='relu', depth=3, dilation_growth_rate=3, down_t=2,
            input_emb_width=140, norm=None, output_emb_width=512,
            type='DecoderAttn', width=512),
        encoder_cfg=dict(
            activation='relu', depth=3, dilation_growth_rate=3, down_t=2,
            input_emb_width=140, norm=None, output_emb_width=512,
            stride_t=2, type='EncoderAttn', width=512),
        quantizer_cfg=dict(dim=512, levels=[8, 8, 6, 5], type='FSQ'),
        type='VQVAE'),
    transformer_cfg=dict(
        block_size=1024, n_embd=512, n_head=8, n_layer=8,
        type='LLaMAHF_Fwd', vocab_size=512),
    type='RetargetTransformerPredMotion_no_smplvq')


# ---------- 推理常量 ----------

FPS = 30
CHUNK_SECONDS = 4
CHUNK_FRAMES = (FPS * CHUNK_SECONDS // 4) * 4  # 120
OVERLAP_FRAMES = 32
STRIDE_FRAMES = CHUNK_FRAMES - OVERLAP_FRAMES


# ---------- 数据加载 ----------

def load_smpl_data(file_path):
    """从 AMASS NPZ 加载 SMPL-X 运动参数。

    支持字段:
      - transl/global_orient/body_pose (标准 SMPL-X)
      - trans/root_orient/pose_body (AMASS 格式，Z-up)
    """
    data = np.load(file_path, allow_pickle=True)

    if 'transl' in data:
        transl = torch.from_numpy(data['transl']).float()
        global_orient = torch.from_numpy(data['global_orient']).float()
        body_pose = torch.from_numpy(data['body_pose']).float()

        if 'mocap_frame_rate' in data:
            src_fps = float(data['mocap_frame_rate'])
            if src_fps > 30:
                step = int(src_fps / 30)
                transl = transl[::step]
                global_orient = global_orient[::step]
                body_pose = body_pose[::step]

        return transl, global_orient, body_pose, None
    else:
        # AMASS 格式：Z-up → Y-up
        transl = torch.from_numpy(data['trans']).float()
        global_orient = torch.from_numpy(data['root_orient']).float()
        body_pose = torch.from_numpy(data['pose_body']).float()

        rot_zup_to_yup = torch.tensor([[1, 0, 0], [0, 0, 1], [0, -1, 0]]).float()
        transl = torch.einsum('ij,tj->ti', rot_zup_to_yup, transl)
        go_mat = axis_angle_to_matrix(global_orient)
        go_mat = torch.einsum('ij,tjk->tik', rot_zup_to_yup, go_mat)
        global_orient = matrix_to_axis_angle(go_mat)

        if 'mocap_frame_rate' in data:
            src_fps = float(data['mocap_frame_rate'])
            if src_fps > 30:
                step = int(src_fps / 30)
                transl = transl[::step]
                global_orient = global_orient[::step]
                body_pose = body_pose[::step]

        return transl, global_orient, body_pose, None


def preprocess_smpl(file_path, smplx_model, betas, device):
    """将 SMPL-X NPZ 转换为 (T, 140) 运动特征向量。"""
    transl, global_orient, body_pose, seq_betas = load_smpl_data(file_path)
    if seq_betas is not None:
        betas = seq_betas

    N = transl.shape[0]
    frame_params = dict(
        transl=transl.to(device),
        global_orient=global_orient.to(device),
        body_pose=body_pose.to(device),
        betas=betas.unsqueeze(0).repeat(N, 1).float().to(device),
        leye_pose=torch.zeros((N, 3), device=device),
        reye_pose=torch.zeros((N, 3), device=device),
        left_hand_pose=torch.zeros((N, 45), device=device),
        right_hand_pose=torch.zeros((N, 45), device=device),
        jaw_pose=torch.zeros((N, 3), device=device),
        expression=torch.zeros((N, 100), device=device),
    )

    with torch.no_grad():
        output = smplx_model(**frame_params)
    position_data = output.joints.detach().cpu()[:, :22]

    global_orient_mat = axis_angle_to_matrix(global_orient)

    position_val_data = position_data[1:] - position_data[:-1]

    root_idx = 0
    y_min = torch.min(position_data[:, :, 1])
    ori = deepcopy(position_data[0, root_idx])
    ori[1] = y_min
    position_data = position_data - ori

    velocities_root = position_data[1:, root_idx, :] - position_data[:-1, root_idx, :]
    position_data_cp = deepcopy(position_data)
    position_data[:, :, 0] -= position_data_cp[:, 0:1, 0]
    position_data[:, :, 2] -= position_data_cp[:, 0:1, 2]

    T, njoint, _ = position_data.shape
    final_x = torch.zeros((T, 2 + 6 + njoint * 3 + njoint * 3))
    final_x[1:, 0] = velocities_root[:, 0]
    final_x[1:, 1] = velocities_root[:, 2]
    final_x[:, 2:8] = matrix_to_rotation_6d(global_orient_mat)
    final_x[:, 8:8 + njoint * 3] = position_data.flatten(1, 2)
    final_x[1:, 8 + njoint * 3:8 + njoint * 6] = position_val_data.flatten(1, 2)

    return final_x


def postprocess_g1(pred_motion, apply_filter=True):
    """从 G1 217 维运动向量提取 DOF/root_trans/root_rot。"""
    T = pred_motion.shape[0]
    rot_mat = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]]).float()

    pred_trans = pred_motion[:, 8:8 + 30 * 3].reshape(T, -1, 3)[:, 0]
    pred_trans[:, [0, 2]] += torch.cumsum(pred_motion[:, :2], dim=0)
    pred_trans = torch.einsum('ij,tj->ti', rot_mat, pred_trans)

    pred_rot_mat = rotation_6d_to_matrix(pred_motion[:, 2:8])
    pred_rot_mat = torch.einsum('ij,tjk->tik', rot_mat, pred_rot_mat)
    pred_rot_quat = matrix_to_quaternion(pred_rot_mat)

    pred_dof = pred_motion[:, -29:]

    if apply_filter and T >= 13:
        _b, _a = sp_signal.butter(4, 5 / (30.0 / 2), btype='low')
        pred_trans = torch.from_numpy(
            sp_signal.filtfilt(_b, _a, pred_trans.numpy(), axis=0).copy()
        ).to(pred_trans.dtype)

        pred_rot_quat_np = sp_signal.filtfilt(_b, _a, pred_rot_quat.numpy(), axis=0).copy()
        pred_rot_quat_np /= np.linalg.norm(pred_rot_quat_np, axis=-1, keepdims=True)
        pred_rot_quat = torch.from_numpy(pred_rot_quat_np).to(pred_rot_quat.dtype)

        pred_dof = torch.from_numpy(
            sp_signal.filtfilt(_b, _a, pred_dof.numpy(), axis=0).copy()
        ).to(pred_dof.dtype)

    return pred_dof, pred_rot_quat, pred_trans


# ---------- 旋转规范化工具 ----------

def _make_y_rot(theta):
    c, s = torch.cos(theta), torch.sin(theta)
    return torch.tensor([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=torch.float32)


def _extract_yaw(rot_6d):
    R = rotation_6d_to_matrix(rot_6d.unsqueeze(0))[0]
    forward = R[:, 2]
    return torch.atan2(forward[0], forward[2])


def _rotate_motion_features(motion, R, n_joints, rotate_6d=True):
    result = motion.clone()
    vx, vz = motion[:, 0], motion[:, 1]
    result[:, 0] = R[0, 0] * vx + R[0, 2] * vz
    result[:, 1] = R[2, 0] * vx + R[2, 2] * vz
    if rotate_6d:
        rot_mat = rotation_6d_to_matrix(motion[:, 2:8])
        rot_mat = torch.einsum('ij,tjk->tik', R, rot_mat)
        result[:, 2:8] = matrix_to_rotation_6d(rot_mat)
    pos_start, pos_end = 8, 8 + n_joints * 3
    pos = motion[:, pos_start:pos_end].reshape(-1, n_joints, 3)
    pos = torch.einsum('ij,tnj->tni', R, pos)
    result[:, pos_start:pos_end] = pos.reshape(-1, n_joints * 3)
    vel_start, vel_end = pos_end, pos_end + n_joints * 3
    if vel_end <= motion.shape[1]:
        vel = motion[:, vel_start:vel_end].reshape(-1, n_joints, 3)
        vel = torch.einsum('ij,tnj->tni', R, vel)
        result[:, vel_start:vel_end] = vel.reshape(-1, n_joints * 3)
    return result


# ---------- 推理核心 ----------

def _infer_chunk(smplx_motion, model, smplx_mean, smplx_std, g1_mean, g1_std, device):
    yaw = _extract_yaw(smplx_motion[0, 2:8])
    R_canon = _make_y_rot(-yaw)
    R_restore = _make_y_rot(yaw)

    smplx_motion = _rotate_motion_features(smplx_motion, R_canon, n_joints=22)

    smplx_motion = (smplx_motion - smplx_mean) / smplx_std
    smplx_input = smplx_motion.unsqueeze(0).float().to(device)
    motion_length = torch.tensor([smplx_motion.shape[0]]).to(device)

    with torch.no_grad():
        pred_motions, _ = model(smplx_motion=smplx_input, motion_length=motion_length, mode='predict')

    pred_motion = pred_motions[0].cpu()
    pred_motion = pred_motion * g1_std + g1_mean

    pred_motion = _rotate_motion_features(pred_motion, R_restore, n_joints=30)
    return pred_motion


def infer_single(file_path, model, smplx_model, betas,
                 smplx_mean, smplx_std, g1_mean, g1_std,
                 device, apply_filter=True):
    t0 = time.time()
    smplx_motion = preprocess_smpl(file_path, smplx_model, betas, device)

    T_orig = smplx_motion.shape[0]
    if T_orig < 4:
        return None, None
    T_pad = ((T_orig + 3) // 4) * 4
    if T_pad > T_orig:
        pad = smplx_motion[-1:].repeat(T_pad - T_orig, 1)
        smplx_motion = torch.cat([smplx_motion, pad], dim=0)
    T = T_pad
    t_preprocess = time.time() - t0

    t1 = time.time()
    if T <= CHUNK_FRAMES:
        pred_motion = _infer_chunk(smplx_motion, model, smplx_mean, smplx_std, g1_mean, g1_std, device)
    else:
        chunks = []
        starts = []
        for start in range(0, T, STRIDE_FRAMES):
            end = min(start + CHUNK_FRAMES, T)
            seg_len = (end - start) // 4 * 4
            if seg_len < 4:
                break
            chunk = smplx_motion[start:start + seg_len]
            chunks.append(_infer_chunk(chunk, model, smplx_mean, smplx_std, g1_mean, g1_std, device))
            starts.append(start)

        pred_motion = chunks[0]
        for i in range(1, len(chunks)):
            overlap = starts[i - 1] + len(chunks[i - 1]) - starts[i]
            if overlap > 0:
                w = torch.linspace(0, 1, overlap).unsqueeze(1)
                prev_tail = pred_motion[-overlap:]
                curr_head = chunks[i][:overlap]
                blended = prev_tail * (1 - w) + curr_head * w
                pred_motion = torch.cat([pred_motion[:-overlap], blended, chunks[i][overlap:]], dim=0)
            else:
                pred_motion = torch.cat([pred_motion, chunks[i]], dim=0)
    t_infer = time.time() - t1

    pred_motion = pred_motion[:T_orig]
    t2 = time.time()
    pred_dof, pred_rot_quat, pred_trans = postprocess_g1(pred_motion, apply_filter=apply_filter)
    t_postprocess = time.time() - t2

    t_total = time.time() - t0
    timing = dict(preprocess=t_preprocess, infer=t_infer, postprocess=t_postprocess, total=t_total)

    return dict(
        dof=pred_dof.numpy(),
        root_trans=pred_trans.numpy(),
        root_rot_quat=pred_rot_quat.numpy(),
        source_path=file_path,
    ), timing


# ---------- 模型加载 ----------

def load_all(weights_dir=None, assets_dir=None, device=None):
    """加载模型、SMPLX 体模型和标准化参数。"""
    if weights_dir is None:
        weights_dir = os.path.join(BASE_DIR, 'weights')
    if assets_dir is None:
        assets_dir = os.path.join(BASE_DIR, 'assets')
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # PyTorch 2.6 兼容
    _torch_load = torch.load
    torch.load = lambda *args, **kwargs: _torch_load(*args, weights_only=kwargs.pop('weights_only', False), **kwargs)

    # 构建并加载模型
    model = MODELS.build(MODEL_CFG)
    ckpt = torch.load(os.path.join(weights_dir, 'epoch_30.pth'), map_location='cpu')
    model.load_state_dict(ckpt['state_dict'])
    model.eval().to(device)

    # 加载 SMPLX 体模型
    smplx_model = SMPLX(
        model_path=os.path.join(assets_dir, 'SMPLX_NEUTRAL.npz'),
        use_pca=False, num_expression_coeffs=100, num_betas=10, ext='npz'
    ).to(device).eval()

    betas = torch.from_numpy(np.load(os.path.join(assets_dir, 'betas.npy'))).float()

    # 标准化参数
    smplx_mean = torch.from_numpy(np.load(os.path.join(weights_dir, 'smplx_mean.npy'))).float()
    smplx_std = torch.from_numpy(np.load(os.path.join(weights_dir, 'smplx_std.npy'))).float()
    g1_mean = torch.from_numpy(np.load(os.path.join(weights_dir, 'gmr_mean.npy'))).float()
    g1_std = torch.from_numpy(np.load(os.path.join(weights_dir, 'gmr_std.npy'))).float()

    return model, smplx_model, betas, smplx_mean, smplx_std, g1_mean, g1_std, device