utils.py

"""
FMPose3D: monocular 3D Pose Estimation via Flow Matching

Official implementation of the paper:
"FMPose3D: monocular 3D Pose Estimation via Flow Matching"
by Ti Wang, Xiaohang Yu, and Mackenzie Weygandt Mathis
Licensed under Apache 2.0
"""

import glob
import hashlib
import json
import os
import shutil

import numpy as np
import torch


def euler_sample(
    c_2d: torch.Tensor,
    y: torch.Tensor,
    steps: int,
    model: torch.nn.Module,
) -> torch.Tensor:
    """Euler ODE sampler for Conditional Flow Matching at test time.

    Integrates the learned velocity field from *t = 0* to *t = 1* using
    ``steps`` uniform Euler steps.

    Parameters
    ----------
    c_2d : Tensor
        2-D conditioning input, shape ``(B, F, J, 2)``.
    y : Tensor
        Initial noise sample (same spatial dims as ``c_2d`` but with 3
        output channels), shape ``(B, F, J, 3)``.
    steps : int
        Number of Euler integration steps.
    model : nn.Module
        The velocity-prediction network ``v(c_2d, y, t)``.

    Returns
    -------
    Tensor
        The denoised 3-D prediction, same shape as *y*.
    """
    dt = 1.0 / steps
    for s in range(steps):
        t_s = torch.full(
            (c_2d.size(0), 1, 1, 1), s * dt, device=c_2d.device, dtype=c_2d.dtype
        )
        v_s = model(c_2d, y, t_s)
        y = y + dt * v_s
    return y

def deterministic_random(min_value, max_value, data):
    digest = hashlib.sha256(data.encode()).digest()
    raw_value = int.from_bytes(digest[:4], byteorder="little", signed=False)
    return int(raw_value / (2**32 - 1) * (max_value - min_value)) + min_value

def mpjpe_cal(predicted, target):
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1))


def test_calculation(predicted, target, action, error_sum, data_type, subject):
    error_sum = mpjpe_by_action_p1(predicted, target, action, error_sum)
    error_sum = mpjpe_by_action_p2(predicted, target, action, error_sum)
    return error_sum


def mpjpe_by_action_p1(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    dist = torch.mean(
        torch.norm(predicted - target, dim=len(target.shape) - 1),
        dim=len(target.shape) - 2,
    )
    if len(set(list(action))) == 1:
        end_index = action[0].find(" ")
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]

        action_error_sum[action_name]["p1"].update(torch.mean(dist).item() * num, num)
    else:
        for i in range(num):
            end_index = action[i].find(" ")
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]

            action_error_sum[action_name]["p1"].update(dist[i].item(), 1)

    return action_error_sum


def p_mpjpe(predicted, target):  # p2, Procrustes analysis MPJPE
    assert predicted.shape == target.shape

    muX = np.mean(target, axis=1, keepdims=True)  # B,1,3
    muY = np.mean(predicted, axis=1, keepdims=True)  # B,1,3

    X0 = target - muX
    Y0 = predicted - muY

    normX = np.sqrt(np.sum(X0**2, axis=(1, 2), keepdims=True))  # B,1,1
    normY = np.sqrt(np.sum(Y0**2, axis=(1, 2), keepdims=True))

    X0 /= normX
    Y0 /= normY

    H = np.matmul(X0.transpose(0, 2, 1), Y0)
    U, s, Vt = np.linalg.svd(H)
    V = Vt.transpose(0, 2, 1)
    R = np.matmul(V, U.transpose(0, 2, 1))

    sign_detR = np.sign(np.expand_dims(np.linalg.det(R), axis=1))
    V[:, :, -1] *= sign_detR
    s[:, -1] *= sign_detR.flatten()
    R = np.matmul(V, U.transpose(0, 2, 1))

    tr = np.expand_dims(np.sum(s, axis=1, keepdims=True), axis=2)

    a = tr * normX / normY
    t = muX - a * np.matmul(muY, R)

    predicted_aligned = a * np.matmul(predicted, R) + t

    return np.mean(
        np.linalg.norm(predicted_aligned - target, axis=len(target.shape) - 1),
        axis=len(target.shape) - 2,
    )


def mpjpe_by_action_p2(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    pred = (
        predicted.detach()
        .cpu()
        .numpy()
        .reshape(-1, predicted.shape[-2], predicted.shape[-1])
    )  # B,17,3
    gt = (
        target.detach().cpu().numpy().reshape(-1, target.shape[-2], target.shape[-1])
    )  # # B,17,3
    dist = p_mpjpe(pred, gt)

    if len(set(list(action))) == 1:
        end_index = action[0].find(" ")
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]
        action_error_sum[action_name]["p2"].update(np.mean(dist) * num, num)
    else:
        for i in range(num):
            end_index = action[i].find(" ")
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]
            action_error_sum[action_name]["p2"].update(np.mean(dist), 1)

    return action_error_sum


def define_actions(action):

    actions = [
        "Directions",
        "Discussion",
        "Eating",
        "Greeting",
        "Phoning",
        "Photo",
        "Posing",
        "Purchases",
        "Sitting",
        "SittingDown",
        "Smoking",
        "Waiting",
        "WalkDog",
        "Walking",
        "WalkTogether",
    ]

    if action == "All" or action == "all" or action == "*":
        return actions

    if not action in actions:
        raise (ValueError, "Unrecognized action: %s" % action)

    return [action]


def define_error_list(actions):
    error_sum = {}
    error_sum.update(
        {
            actions[i]: {"p1": AccumLoss(), "p2": AccumLoss()}
            for i in range(len(actions))
        }
    )
    return error_sum


class AccumLoss(object):
    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val
        self.count += n
        self.avg = self.sum / self.count


def get_variable(split, target):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num = len(target)
    var = []
    if split == "train":
        for i in range(num):
            temp = target[i].requires_grad_(False).contiguous().float().to(device)
            var.append(temp)
    else:
        for i in range(num):
            temp = target[i].contiguous().float().to(device)
            var.append(temp)

    return var


def print_error(data_type, action_error_sum, is_train):
    mean_error_p1, mean_error_p2 = print_error_action(action_error_sum, is_train)

    return mean_error_p1, mean_error_p2


def print_error_action(action_error_sum, is_train):
    mean_error_each = {"p1": 0.0, "p2": 0.0}
    mean_error_all = {"p1": AccumLoss(), "p2": AccumLoss()}

    if is_train == 0:
        print("{0:=^12} {1:=^10} {2:=^8}".format("Action", "p#1 mm", "p#2 mm"))

    for action, value in action_error_sum.items():
        if is_train == 0:
            print("{0:<12} ".format(action), end="")

        mean_error_each["p1"] = action_error_sum[action]["p1"].avg * 1000.0
        mean_error_all["p1"].update(mean_error_each["p1"], 1)

        mean_error_each["p2"] = action_error_sum[action]["p2"].avg * 1000.0
        mean_error_all["p2"].update(mean_error_each["p2"], 1)

        if is_train == 0:
            print(
                "{0:>6.2f} {1:>10.2f}".format(
                    mean_error_each["p1"], mean_error_each["p2"]
                )
            )

    if is_train == 0:
        print(
            "{0:<12} {1:>6.4f} {2:>10.4f}".format(
                "Average", mean_error_all["p1"].avg, mean_error_all["p2"].avg
            )
        )

    return mean_error_all["p1"].avg, mean_error_all["p2"].avg


def save_model(previous_name, save_dir, epoch, data_threshold, model, model_name):
    # remove the old model
    if os.path.exists(previous_name):
        os.remove(previous_name)

    torch.save(
        model.state_dict(),
        "%s/%s_%d_%d.pth" % (save_dir, model_name, epoch, data_threshold * 100),
    )
    previous_name = "%s/%s_%d_%d.pth" % (
        save_dir,
        model_name,
        epoch,
        data_threshold * 100,
    )
    return previous_name


def save_top_N_models(
    previous_name,
    save_dir,
    epoch,
    data_threshold,
    model,
    model_name,
    num_saved_models=3,
):
    """
    Save a checkpoint if it belongs to the top-N best (by lower data_threshold).
    Maintains an index file '<model_name>_top_models.json' in save_dir.

    Returns the path of the last saved checkpoint if a new one was saved,
    otherwise returns previous_name unchanged.
    """
    os.makedirs(save_dir, exist_ok=True)
    ckpt_path = os.path.join(
        save_dir, f"{model_name}_{epoch}_{int(data_threshold * 100)}.pth"
    )
    index_path = os.path.join(save_dir, f"{model_name}_top_models.json")

    # load current list
    top_list = []
    if os.path.exists(index_path):
        with open(index_path, "r") as f:
            top_list = json.load(f)

    # decide if we should save
    should_save = False
    if len(top_list) < int(num_saved_models):
        should_save = True
    else:
        worst_item = max(top_list, key=lambda x: x.get("p1", float("inf")))
        if data_threshold < float(worst_item.get("p1", float("inf"))):
            should_save = True

    if not should_save:
        return previous_name

    # save new checkpoint
    torch.save(model.state_dict(), ckpt_path)

    # append and trim to N
    top_list.append(
        {"p1": float(data_threshold), "path": ckpt_path, "epoch": int(epoch)}
    )
    # sort ascending by p1 and keep best N
    top_list.sort(key=lambda x: x.get("p1", float("inf")))
    while len(top_list) > int(num_saved_models):
        removed = top_list.pop()  # remove worst (last after sort ascending)
        if os.path.exists(removed.get("path", "")):
            os.remove(removed["path"])

    # write back index
    with open(index_path, "w") as f:
        json.dump(top_list, f, indent=2)

    # update best marker to point to current best (lowest p1): append _best to the original name
    if len(top_list) > 0:
        best_src = top_list[0].get("path")
        if best_src and os.path.exists(best_src):
            root_name, ext = os.path.splitext(best_src)
            best_path = f"{root_name}_best{ext}"
            try:
                # ensure only one best exists: remove all existing *_best for this model_name
                pattern = os.path.join(save_dir, f"{model_name}_*_best.pth")
                for old_best in glob.glob(pattern):
                    try:
                        os.remove(old_best)
                    except Exception:
                        pass
                shutil.copy2(best_src, best_path)
            except Exception:
                pass

    return ckpt_path


def back_to_ori_uv(cropped_uv, bb_box):
    """
    for cropped uv, back to original uv to help do the uvd->xyz operation
    :return:
    """
    N, T, V, _ = cropped_uv.size()
    uv = (cropped_uv + 1) * (bb_box[:, 2:].view(N, 1, 1, 2) / 2.0) + bb_box[
        :, 0:2
    ].view(N, 1, 1, 2)
    return uv


def get_uvd2xyz(uvd, gt_3D, cam):
    """
    transfer uvd to xyz

    :param uvd: N*T*V*3 (uv and z channel)
    :param gt_3D: N*T*V*3 (NOTE: V=0 is absolute depth value of root joint)

    :return: root-relative xyz results
    """
    N, T, V, _ = uvd.size()

    dec_out_all = uvd.view(-1, T, V, 3).clone()  # N*T*V*3
    root = gt_3D[:, :, 0, :].unsqueeze(-2).repeat(1, 1, V, 1).clone()  # N*T*V*3
    enc_in_all = uvd[:, :, :, :2].view(-1, T, V, 2).clone()  # N*T*V*2

    cam_f_all = cam[..., :2].view(-1, 1, 1, 2).repeat(1, T, V, 1)  # N*T*V*2
    cam_c_all = cam[..., 2:4].view(-1, 1, 1, 2).repeat(1, T, V, 1)  # N*T*V*2

    # change to global
    z_global = dec_out_all[:, :, :, 2]  # N*T*V
    z_global[:, :, 0] = root[:, :, 0, 2]
    z_global[:, :, 1:] = dec_out_all[:, :, 1:, 2] + root[:, :, 1:, 2]  # N*T*V
    z_global = z_global.unsqueeze(-1)  # N*T*V*1

    uv = enc_in_all - cam_c_all  # N*T*V*2
    xy = uv * z_global.repeat(1, 1, 1, 2) / cam_f_all  # N*T*V*2
    xyz_global = torch.cat((xy, z_global), -1)  # N*T*V*3
    xyz_offset = xyz_global - xyz_global[:, :, 0, :].unsqueeze(-2).repeat(
        1, 1, V, 1
    )  # N*T*V*3

    return xyz_offset


def sym_penalty(dataset, keypoints, pred_out):
    """
    get penalty for the symmetry of human body
    :return:
    """
    loss_sym = 0
    if dataset == "h36m":
        if keypoints.startswith("sh"):
            left_bone = [(0, 4), (4, 5), (5, 6), (8, 10), (10, 11), (11, 12)]
            right_bone = [(0, 1), (1, 2), (2, 3), (8, 13), (13, 14), (14, 15)]
        else:
            left_bone = [(0, 4), (4, 5), (5, 6), (8, 11), (11, 12), (12, 13)]
            right_bone = [(0, 1), (1, 2), (2, 3), (8, 14), (14, 15), (15, 16)]
        for (i_left, j_left), (i_right, j_right) in zip(left_bone, right_bone):
            left_part = pred_out[:, :, i_left] - pred_out[:, :, j_left]
            right_part = pred_out[:, :, i_right] - pred_out[:, :, j_right]
            loss_sym += torch.mean(
                torch.norm(left_part, dim=-1) - torch.norm(right_part, dim=-1)
            )
    elif dataset.startswith("STB"):
        loss_sym = 0
    return loss_sym


def project_to_2d(X, camera_params):
    """
    Project 3D points to 2D using the Human3.6M camera projection function.
    This is a differentiable and batched reimplementation of the original MATLAB script.

    Arguments:
    X -- 3D points in *camera space* to transform (N, *, 3)
    camera_params -- intrinsic parameters (N, 2+2+3+2=9)
    """
    assert X.shape[-1] == 3  #  B,J,3
    assert len(camera_params.shape) == 2  # camera_params:[B,1,9]
    assert camera_params.shape[-1] == 9
    assert X.shape[0] == camera_params.shape[0]

    while len(camera_params.shape) < len(X.shape):
        camera_params = camera_params.unsqueeze(1)

    f = camera_params[..., :2]
    c = camera_params[..., 2:4]
    k = camera_params[..., 4:7]  # B,1,3
    p = camera_params[..., 7:]  # B,1,2

    XX = torch.clamp(X[..., :2] / X[..., 2:], min=-1, max=1)  # B,J,2
    r2 = torch.sum(XX[..., :2] ** 2, dim=len(XX.shape) - 1, keepdim=True)  # B, J, 1

    radial = 1 + torch.sum(
        k * torch.cat((r2, r2**2, r2**3), dim=len(r2.shape) - 1),
        dim=len(r2.shape) - 1,  # B,J,1
        keepdim=True,
    )
    tan = torch.sum(p * XX, dim=len(XX.shape) - 1, keepdim=True)  # B,J,1

    XXX = XX * (radial + tan) + p * r2  # B,J,2

    return f * XXX + c


def input_augmentation(input_2D, model):
    joints_left = [4, 5, 6, 11, 12, 13]
    joints_right = [1, 2, 3, 14, 15, 16]

    input_2D_non_flip = input_2D[:, 0]
    input_2D_flip = input_2D[:, 1]

    output_3D_non_flip = model(input_2D_non_flip)
    output_3D_flip = model(input_2D_flip)

    output_3D_flip[:, :, :, 0] *= -1
    output_3D_flip[:, :, joints_left + joints_right, :] = output_3D_flip[
        :, :, joints_right + joints_left, :
    ]

    output_3D = (output_3D_non_flip + output_3D_flip) / 2

    input_2D = input_2D_non_flip

    return input_2D, output_3D