fmpose3d 0.0.7__py3-none-any.whl

fmpose3d/__init__.py

@@ -0,0 +1,38 @@
+"""
+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
+"""
+
+__version__ = "0.0.7"
+__author__ = "Ti Wang, Xiaohang Yu, Mackenzie Weygandt Mathis"
+__license__ = "Apache 2.0"
+
+# Import key components for easy access
+from .aggregation_methods import (
+    average_aggregation,
+    aggregation_select_single_best_hypothesis_by_2D_error,
+    aggregation_RPEA_joint_level,
+)
+
+# Import 2D pose detection utilities
+from .lib.hrnet.gen_kpts import gen_video_kpts
+from .lib.preprocess import h36m_coco_format, revise_kpts
+
+# Make commonly used classes/functions available at package level
+__all__ = [
+    # Aggregation methods
+    "average_aggregation",
+    "aggregation_select_single_best_hypothesis_by_2D_error",
+    "aggregation_RPEA_joint_level",
+    # 2D pose detection
+    "gen_video_kpts",
+    "h36m_coco_format",
+    "revise_kpts",
+    # Version
+    "__version__",
+]
+

fmpose3d/aggregation_methods.py

@@ -0,0 +1,203 @@
+"""
+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 torch
+from fmpose3d.common.utils import project_to_2d
+
+def average_aggregation(list_hypothesis):
+    return torch.mean(torch.stack(list_hypothesis), dim=0)
+
+
+def aggregation_select_single_best_hypothesis_by_2D_error(args,
+    list_hypothesis, batch_cam, input_2D, gt_3D
+):
+    """
+    Select per-joint 3D from the hypothesis whose 2D projection yields minimal L2 error.
+
+    Args:
+        list_hypothesis: list of (B,1,J,3) tensors
+        batch_cam: (B, 9) or (B, 1, 9) intrinsics [f(2), c(2), k(3), p(2)]
+        input_2D: (B, F, J, 2) 2D joints in image coordinates
+        gt_3D: (B, F, J, 3) used for shape metadata only
+    Returns:
+        (B,1,J,3) aggregated 3D pose with joint 0 set to 0
+    """
+    if len(list_hypothesis) == 0:
+        raise ValueError("list_hypothesis is empty")
+
+    device = list_hypothesis[0].device
+    dtype = list_hypothesis[0].dtype
+
+    # Shapes
+    B = gt_3D.size(0)
+    J = gt_3D.size(2)
+    F = gt_3D.size(1)
+    assert F >= 1, "Expected at least one frame"
+
+    # Stack hypotheses: (H,B,1,J,3) -> (B,H,J,3)
+    stack = torch.stack(list_hypothesis, dim=0)  # (H,B,1,J,3)
+    X_hbj3 = stack[:, :, 0, :, :]  # (H,B,J,3)
+    X_bhj3 = X_hbj3.transpose(0, 1).contiguous()  # (B,H,J,3)
+    H = X_bhj3.size(1)
+
+    # Prepare camera params: (B,9)
+    if batch_cam.dim() == 3 and batch_cam.size(1) == 1:
+        cam_b9 = batch_cam[:, 0, :].contiguous()
+    elif batch_cam.dim() == 2 and batch_cam.size(1) == 9:
+        cam_b9 = batch_cam
+    else:
+        cam_b9 = batch_cam.view(B, -1)
+    assert cam_b9.size(1) == 9, f"camera params should be 9-dim, got {cam_b9.size()}"
+
+    # Target 2D at the same frame index as 3D selection (args.pad)
+    # input_2D: (B,F,J,2) -> (B,J,2)
+    target_2d = input_2D[:, getattr(args, "pad", 0)].contiguous()  # (B,J,2)
+
+    # Convert hypotheses from root-relative to absolute camera coordinates using GT root
+    # Root at frame args.pad: (B,3)
+    gt_root = gt_3D[:, getattr(args, "pad", 0), 0, :].contiguous()  # (B,3)
+    X_abs = X_bhj3.clone()
+    X_abs[:, :, 1:, :] = X_abs[:, :, 1:, :] + gt_root.unsqueeze(1).unsqueeze(1)
+    X_abs[:, :, 0, :] = gt_root.unsqueeze(1)
+
+    # Vectorized projection for all hypotheses in absolute coordinates
+    # (B,H,J,3) -> (B*H,J,3)
+    X_flat = X_abs.view(B * H, J, 3)
+    cam_rep = cam_b9.repeat_interleave(H, dim=0)  # (B*H,9)
+
+    # project_to_2d expects last dim=3 and cam (N,9)
+    # Returns normalized coordinates (when crop_uv=0) because camera params are normalized
+    proj2d_flat = project_to_2d(X_flat, cam_rep)  # (B*H,J,2) normalized coordinates
+    proj2d_bhj = proj2d_flat.view(B, H, J, 2)
+
+    # Per-hypothesis per-joint 2D error (both in normalized coordinates)
+    diff = proj2d_bhj - target_2d.unsqueeze(1)  # (B,H,J,2)
+    dist = torch.norm(diff, dim=-1)  # (B,H,J)
+
+    # Exclude root joint (0) due to undefined depth when using root-relative 3D
+    dist[:, :, 0] = float("inf")
+
+    # Argmin across hypotheses per joint
+    best_h = torch.argmin(dist, dim=1)  # (B,J)
+
+    # Gather 3D using advanced indexing (return root-relative coordinates)
+    b_idx = torch.arange(B, device=device).unsqueeze(1).expand(B, J)  # (B,J)
+    j_idx = torch.arange(J, device=device).unsqueeze(0).expand(B, J)  # (B,J)
+    selected_bj3 = X_bhj3[b_idx, best_h, j_idx, :]  # (B,J,3)
+
+    agg = selected_bj3.unsqueeze(1).to(dtype=dtype)
+    agg[:, :, 0, :] = 0
+    return agg
+
+
+def aggregation_RPEA_joint_level(
+    args, list_hypothesis, batch_cam, input_2D, gt_3D, topk=3
+):
+    """
+    Select per-joint 3D from the hypothesis whose 2D projection yields minimal L2 error.
+
+    Args:
+        list_hypothesis: list of (B,1,J,3) tensors
+        batch_cam: (B, 9) or (B, 1, 9) intrinsics [f(2), c(2), k(3), p(2)]
+        input_2D: (B, F, J, 2) 2D joints in image coordinates
+        gt_3D: (B, F, J, 3) used for shape metadata only
+    Returns:
+        (B,1,J,3) aggregated 3D pose with joint 0 set to 0
+    """
+    if len(list_hypothesis) == 0:
+        raise ValueError("list_hypothesis is empty")
+
+    device = list_hypothesis[0].device
+    dtype = list_hypothesis[0].dtype
+
+    # Shapes
+    B = gt_3D.size(0)
+    J = gt_3D.size(2)
+    F = gt_3D.size(1)
+    assert F >= 1, "Expected at least one frame"
+
+    # Stack hypotheses: (H,B,1,J,3) -> (B,H,J,3)
+    stack = torch.stack(list_hypothesis, dim=0)  # (H,B,1,J,3)
+    X_hbj3 = stack[:, :, 0, :, :]  # (H,B,J,3)
+    X_bhj3 = X_hbj3.transpose(0, 1).contiguous()  # (B,H,J,3)
+    H = X_bhj3.size(1)
+
+    # Prepare camera params: (B,9)
+    if batch_cam.dim() == 3 and batch_cam.size(1) == 1:
+        cam_b9 = batch_cam[:, 0, :].contiguous()
+    elif batch_cam.dim() == 2 and batch_cam.size(1) == 9:
+        cam_b9 = batch_cam
+    else:
+        cam_b9 = batch_cam.view(B, -1)
+    assert cam_b9.size(1) == 9, f"camera params should be 9-dim, got {cam_b9.size()}"
+
+    # Target 2D at the same frame index as 3D selection (args.pad)
+    # input_2D: (B,F,J,2) -> (B,J,2)
+    target_2d = input_2D[:, getattr(args, "pad", 0)].contiguous()  # (B,J,2)
+
+    # Convert hypotheses from root-relative to absolute camera coordinates using GT root
+    # Root at frame args.pad: (B,3)
+    gt_root = gt_3D[:, getattr(args, "pad", 0), 0, :].contiguous()  # (B,3)
+    X_abs = X_bhj3.clone()
+    X_abs[:, :, 1:, :] = X_abs[:, :, 1:, :] + gt_root.unsqueeze(1).unsqueeze(1)
+    X_abs[:, :, 0, :] = gt_root.unsqueeze(1)
+
+    # Vectorized projection for all hypotheses in absolute coordinates
+    # (B,H,J,3) -> (B*H,J,3)
+    X_flat = X_abs.view(B * H, J, 3)
+    cam_rep = cam_b9.repeat_interleave(H, dim=0)  # (B*H,9)
+
+    # project_to_2d expects last dim=3 and cam (N,9)
+    proj2d_flat = project_to_2d(X_flat, cam_rep)  # (B*H,J,2)
+    proj2d_bhj = proj2d_flat.view(B, H, J, 2)
+
+    # Per-hypothesis per-joint 2D error
+    diff = proj2d_bhj - target_2d.unsqueeze(1)  # (B,H,J,2)
+    dist = torch.norm(diff, dim=-1)  # (B,H,J)
+
+    # For root joint (0), avoid NaNs in softmax by setting equal distances
+    # This yields uniform weights at the root (we set root to 0 later anyway)
+    dist[:, :, 0] = 0.0
+
+    # Convert 2D losses to weights using softmax over top-k hypotheses per joint
+    tau = float(getattr(args, "weight_softmax_tau", 1.0))
+    H = dist.size(1)
+    k = int(getattr(args, "topk", None))
+    # print("k:", k)
+    # k = int(H//2)+1
+    k = max(1, min(k, H))
+
+    # top-k smallest distances along hypothesis dim
+    topk_vals, topk_idx = torch.topk(dist, k=k, dim=1, largest=False)  # (B,k,J)
+
+    # Weight calculation method ; weight_method = 'exp'
+    temp = args.exp_temp
+    max_safe_val = temp * 20
+    topk_vals_clipped = torch.clamp(topk_vals, max=max_safe_val)
+    exp_vals = torch.exp(-topk_vals_clipped / temp)
+    exp_sum = exp_vals.sum(dim=1, keepdim=True)
+    topk_weights = exp_vals / torch.clamp(exp_sum, min=1e-10)
+    nan_mask = torch.isnan(topk_weights).any(dim=1, keepdim=True)
+    uniform_weights = torch.ones_like(topk_weights) / k
+    topk_weights = torch.where(
+        nan_mask.expand_as(topk_weights), uniform_weights, topk_weights
+    )
+
+    # scatter back to full H with zeros elsewhere
+    weights = torch.zeros_like(dist)  # (B,H,J)
+    weights.scatter_(1, topk_idx, topk_weights)
+
+    # Weighted sum of root-relative 3D hypotheses per joint
+    weights_exp = weights.unsqueeze(-1)  # (B,H,J,1)
+    weighted_bj3 = torch.sum(X_bhj3 * weights_exp, dim=1)  # (B,J,3)
+
+    # Assemble output (B,1,J,3) and enforce root at origin
+    agg = weighted_bj3.unsqueeze(1).to(dtype=dtype)
+    agg[:, :, 0, :] = 0
+    return agg

fmpose3d/animals/__init__.py

@@ -0,0 +1,13 @@
+"""
+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
+"""
+
+__all__ = [
+    "common",
+]
+

fmpose3d/animals/common/__init__.py

@@ -0,0 +1,12 @@
+"""
+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
+"""
+
+"""
+Shared utilities for animal datasets and models.
+"""

fmpose3d/animals/common/animal3d_dataset.py

@@ -0,0 +1,66 @@
+"""
+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 json
+
+import numpy as np
+from .camera import normalize_screen_coordinates
+from torch.utils.data import Dataset
+
+
+class TrainDataset(Dataset):
+    def __init__(self, is_train: bool, json_file: str, root_joint: int = 12):
+        super().__init__()
+        self.focal_length = 1000
+        self.root_joint = root_joint  # Root joint index for making coordinates relative
+
+        json_file = json_file
+        with open(json_file, "r") as f:
+            self.data = json.load(f)
+
+        self.is_train = is_train
+
+    def __len__(self):
+        return len(self.data["data"])
+
+    def __getitem__(self, item):
+        data = self.data["data"][item]
+        # safely check for reproj_kp_2d
+        reproj = data.get("reproj_kp_2d", None)
+        if reproj is not None:
+            keypoint_2d = np.array(reproj, dtype=np.float32)
+        else:
+            keypoint_2d = np.array(data.get("keypoint_2d", []), dtype=np.float32)
+         # normalize 2D keypoints
+        hight = np.array(data["height"])
+        width = np.array(data["width"])
+        keypoint_2d = normalize_screen_coordinates(keypoint_2d[..., :2], width, hight)
+         
+        # build 3D keypoints; append ones; fallback to zeros if missing
+        if "keypoint_3d" in data and data["keypoint_3d"] is not None:
+            kp3d = np.array(data["keypoint_3d"], dtype=np.float32)
+            keypoint_3d = np.concatenate((kp3d, np.ones((len(kp3d), 1), dtype=np.float32)), axis=-1)
+        else:
+            keypoint_3d = np.zeros((len(keypoint_2d), 4), dtype=np.float32)
+
+        # Make 3D keypoints root-relative
+        if keypoint_3d.shape[0] > self.root_joint:
+            root_pos = keypoint_3d[self.root_joint : self.root_joint + 1, :].copy()  # (1, 4)
+            keypoint_3d = keypoint_3d - root_pos  # All joints relative to root
+            # Now root joint should be exactly [0,0,0,0]
+
+        bbox = data["bbox"]  # [x, y, w, h]
+        ori_keypoint_2d = keypoint_2d.copy()
+  
+        item = {
+            "keypoints_2d": keypoint_2d,  #
+            "keypoints_3d": keypoint_3d,
+            "img_path": str(data["img_path"]),
+              }
+        return item

fmpose3d/animals/common/animal_visualization.py

@@ -0,0 +1,111 @@
+"""
+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 os
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+def save_3Dpose_colored(pre_pose, gt_pose, figure_name):
+    fig = plt.figure()
+    ax1 = fig.add_subplot(211, projection="3d")
+    ax1.scatter(
+        pre_pose[:, 0], pre_pose[:, 1], pre_pose[:, 2], c=list(range(pre_pose.shape[0])), cmap="jet"
+    )
+    # plt.axis('off')
+    ax2 = fig.add_subplot(212, projection="3d")
+    ax2.scatter(
+        gt_pose[:, 0], gt_pose[:, 1], gt_pose[:, 2], c=list(range(gt_pose.shape[0])), cmap="jet"
+    )
+    # plt.axis('off')
+    plt.show()
+    plt.savefig(figure_name, dpi=400.0)
+    plt.close()
+
+
+def save_absolute_3Dpose_image(image, pre_pose, gt_pose, vid_3D, skeleton, figure_name):
+    fig = plt.figure(figsize=(20, 9))
+    ax1 = fig.add_subplot(131, projection="3d")
+    ax1.scatter(
+        pre_pose[:, 0],
+        pre_pose[:, 2],
+        -pre_pose[:, 1],
+        c=list(range(pre_pose.shape[0])),
+        cmap="jet",
+    )
+    for i in range(skeleton.shape[0]):
+        ax1.plot(
+            [pre_pose[skeleton[i, 0], 0], pre_pose[skeleton[i, 1], 0]],
+            [pre_pose[skeleton[i, 0], 2], pre_pose[skeleton[i, 1], 2]],
+            [-pre_pose[skeleton[i, 0], 1], -pre_pose[skeleton[i, 1], 1]],
+            c="black",
+        )
+    ax1.set_xlim([-3, 3])
+    ax1.set_zlim([-1.5, 3])
+    ax1.set_ylim([12, 20])
+    ax1.title.set_text("Prediction")
+
+    # plt.axis('off')
+    ax2 = fig.add_subplot(132, projection="3d")
+    visiable_gt = gt_pose[np.where(vid_3D)[0], :]
+    ax2.scatter(
+        visiable_gt[:, 0],
+        visiable_gt[:, 2],
+        -visiable_gt[:, 1],
+        c=list(np.array(range(gt_pose.shape[0]))[np.where(vid_3D)]),
+        cmap="jet",
+    )
+    for i in range(skeleton.shape[0]):
+        if vid_3D[skeleton[i, 0]] > 0 and vid_3D[skeleton[i, 1]] > 0:
+            ax2.plot(
+                [gt_pose[skeleton[i, 0], 0], gt_pose[skeleton[i, 1], 0]],
+                [gt_pose[skeleton[i, 0], 2], gt_pose[skeleton[i, 1], 2]],
+                [-gt_pose[skeleton[i, 0], 1], -gt_pose[skeleton[i, 1], 1]],
+                c="black",
+            )
+    ax2.set_xlim([-3, 3])
+    ax2.set_zlim([-1.5, 3])
+    ax2.set_ylim([12, 20])
+    ax2.title.set_text("GT")
+    # plt.axis('off')
+    ax3 = fig.add_subplot(133)
+    ax3.imshow(image)
+    ax3.title.set_text("Camera1 view")
+    plt.show()
+    plt.savefig(figure_name, dpi=200.0)
+    plt.close()
+
+
+def save_absolute_3Dpose(pre_pose, skeleton, figure_name):
+    fig = plt.figure(figsize=(20, 9))
+    ax1 = fig.add_subplot(111, projection="3d")
+    ax1.scatter(
+        pre_pose[:, 0],
+        pre_pose[:, 2],
+        -pre_pose[:, 1],
+        c=list(range(pre_pose.shape[0])),
+        cmap="jet",
+    )
+    for i in range(skeleton.shape[0]):
+        ax1.plot(
+            [pre_pose[skeleton[i, 0], 0], pre_pose[skeleton[i, 1], 0]],
+            [pre_pose[skeleton[i, 0], 2], pre_pose[skeleton[i, 1], 2]],
+            [-pre_pose[skeleton[i, 0], 1], -pre_pose[skeleton[i, 1], 1]],
+            c="black",
+        )
+    ax1.set_xlim([-1, 1])
+    ax1.set_zlim([-1, 1])
+    ax1.set_ylim([-1, 1])
+    ax1.title.set_text("gt")
+    plt.show()
+    plt.savefig(figure_name, dpi=200.0)
+    plt.close()