diff --git a/README.md b/README.md index a77c5859..e56521d3 100644 --- a/README.md +++ b/README.md @@ -1,17 +1,25 @@ -# FMPose: 3D Pose Estimation via Flow Matching +# FMPose3D: monocular 3D Pose Estimation via Flow Matching This is the official implementation of the approach described in the paper: -> [**FMPose: 3D Pose Estimation via Flow Matching**](xxx) +> [**FMPose3D: monocular 3D Pose Estimation via Flow Matching**](xxx) > Ti Wang, Xiaohang Yu, Mackenzie Weygandt Mathis

-## Set up a environment -Make sure you have Python 3.10. You can set this up with: +## News! + +- [X] Feb 2026: FMPose3D is code and arXiv paper is released - check out the demos here or on our [project page](https://xiu-cs.github.io/FMPose3D/) +- [ ] Planned: This method will be integrated into [DeepLabCut](https://www.mackenziemathislab.org/deeplabcut) + +## Installation + +### Set up an environment + +Make sure you have Python 3.10+. You can set this up with: ```bash conda create -n fmpose_3d python=3.10 conda activate fmpose_3d @@ -20,9 +28,9 @@ conda activate fmpose_3d ```bash git clone xxxx.git # clone this repo # TestPyPI (pre-release/testing build) -pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ fmpose==0.0.5 +pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ fmpose3d==0.0.7 # Future Official PyPI release -# pip install fmpose +# pip install fmpose3d ``` ## Demo @@ -68,7 +76,7 @@ The training logs, checkpoints, and related files of each training time will be For training on Human3.6M: ```bash -sh /scripts/FMPose_train.sh +sh /scripts/FMPose3D_train.sh ``` ### Inference @@ -78,16 +86,18 @@ First, download the folder with pre-trained model from [here](https://drive.goog To run inference on Human3.6M: ```bash -sh ./scripts/FMPose_test.sh +sh ./scripts/FMPose3D_test.sh ``` ## Experiments Animals For animal training/testing and demo scripts, see [animals/README.md](animals/README.md). -## Acknowledgement +## Acknowledgements + +We thank the Swiss National Science Foundation (SNSF Project # 320030-227871) and the Kavli Foundation for providing financial support for this project. -Our code is extended from the following repositories. We thank the authors for releasing the codes. +Our code is extended from the following repositories. We thank the authors for releasing the code. - [MHFormer](https://github.com/Vegetebird/MHFormer) - [StridedTransformer-Pose3D](https://github.com/Vegetebird/StridedTransformer-Pose3D) diff --git a/animals/README.md b/animals/README.md index 49bd922f..af9da699 100644 --- a/animals/README.md +++ b/animals/README.md @@ -2,7 +2,7 @@ # Animals -In this part, the FMPose model is trained on [Animal3D](https://xujiacong.github.io/Animal3D/) dataset and [Control_Animal3D](https://luoxue-star.github.io/AniMer_project_page/) dataset. +In this part, the FMPose3D model is trained on [Animal3D](https://xujiacong.github.io/Animal3D/) dataset and [Control_Animal3D](https://luoxue-star.github.io/AniMer_project_page/) dataset. ## Demo ### Testing on in-the-wild images (animals) @@ -62,4 +62,4 @@ Download the pretrained model from [here](https://drive.google.com/drive/folders ```bash cd animals # the current path is: ./animals bash ./scripts/test_animal3d.sh -``` \ No newline at end of file +``` diff --git a/animals/demo/vis_animals.py b/animals/demo/vis_animals.py index 204c049d..357cfe80 100644 --- a/animals/demo/vis_animals.py +++ b/animals/demo/vis_animals.py @@ -22,8 +22,8 @@ from PIL import Image import matplotlib.gridspec as gridspec import imageio -from fmpose.animals.common.arguments import opts as parse_args -from fmpose.common.camera import normalize_screen_coordinates, camera_to_world +from fmpose3d.animals.common.arguments import opts as parse_args +from fmpose3d.common.camera import normalize_screen_coordinates, camera_to_world sys.path.append(os.getcwd()) @@ -46,7 +46,7 @@ CFM = getattr(module, "Model") else: # Load model from installed fmpose package - from fmpose.models import Model as CFM + from fmpose3d.models import Model as CFM from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images diff --git a/animals/demo/vis_animals.sh b/animals/demo/vis_animals.sh index 66a7f0b9..38793919 100644 --- a/animals/demo/vis_animals.sh +++ b/animals/demo/vis_animals.sh @@ -7,7 +7,7 @@ sh_file='vis_animals.sh' # n_joints=26 # out_joints=26 -model_path='../../fmpose/animals/models/model_animal3d.py' +model_path='../pre_trained_models/animal3d_pretrained_weights/model_animal3d.py' saved_model_path='../pre_trained_models/animal3d_pretrained_weights/CFM_154_4403_best.pth' # path='./images/image_00068.jpg' # single image diff --git a/animals/models/model_animals.py b/animals/models/model_animals.py index a1f95d2f..4d872a68 100644 --- a/animals/models/model_animals.py +++ b/animals/models/model_animals.py @@ -15,7 +15,7 @@ from einops import rearrange from timm.models.layers import DropPath -from fmpose.animals.models.graph_frames import Graph +from fmpose3d.animals.models.graph_frames import Graph class TimeEmbedding(nn.Module): @@ -148,7 +148,7 @@ def forward(self, x): x = res2 + self.drop_path(x) return x -class FMPose(nn.Module): +class FMPose3D(nn.Module): def __init__(self, depth, embed_dim, channels_dim, tokens_dim, adj, drop_rate=0.10, length=27): super().__init__() drop_path_rate = 0.2 @@ -220,7 +220,7 @@ def __init__(self, args): self.encoder_pose_2d = encoder(2, args.channel//2, args.channel//2-self.t_embed_dim//2) self.encoder_y_t = encoder(3, args.channel//2, args.channel//2-self.t_embed_dim//2) - self.FMPose = FMPose(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256 + self.FMPose3D = FMPose3D(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256 self.pred_mu = decoder(args.channel, args.channel//2, 3) def forward(self, pose_2d, y_t, t): @@ -239,7 +239,7 @@ def forward(self, pose_2d, y_t, t): in_emb = rearrange(in_emb, 'b f j c -> (b f) j c').contiguous() # (B*F,J,in) # encoder -> model -> regression head - h = self.FMPose(in_emb) + h = self.FMPose3D(in_emb) v = self.pred_mu(h) # (B*F,J,3) v = rearrange(v, '(b f) j c -> b f j c', b=b, f=f).contiguous() # (B,F,J,3) diff --git a/animals/scripts/main_animal3d.py b/animals/scripts/main_animal3d.py index 157b5d52..f93e04cf 100644 --- a/animals/scripts/main_animal3d.py +++ b/animals/scripts/main_animal3d.py @@ -15,9 +15,9 @@ import numpy as np from tqdm import tqdm import torch.optim as optim -from fmpose.animals.common.arguments import opts as parse_args -from fmpose.animals.common.utils import * -from fmpose.animals.common.animal3d_dataset import TrainDataset +from fmpose3d.animals.common.arguments import opts as parse_args +from fmpose3d.animals.common.utils import * +from fmpose3d.animals.common.animal3d_dataset import TrainDataset import time args = parse_args().parse() @@ -39,7 +39,7 @@ CFM = getattr(module, "Model") else: # Load model from installed fmpose package - from fmpose.animals.models import Model as CFM + from fmpose3d.animals.models import Model as CFM def train(opt, actions, train_loader, model, optimizer, epoch): return step('train', opt, actions, train_loader, model, optimizer, epoch) diff --git a/animals/scripts/test_animal3d.sh b/animals/scripts/test_animal3d.sh index 65b61779..3c858854 100644 --- a/animals/scripts/test_animal3d.sh +++ b/animals/scripts/test_animal3d.sh @@ -11,7 +11,7 @@ n_joints=26 out_joints=26 epochs=300 # model_path='models/model_animals.py' -model_path="" # when the path is empty, the model will be loaded from the installed fmpose package +model_path='./pre_trained_models/animal3d_pretrained_weights/model_animal3d.py' # when the path is empty, the model will be loaded from the installed fmpose package saved_model_path='./pre_trained_models/animal3d_pretrained_weights/CFM_154_4403_best.pth' # root path denotes the path to the original dataset root_path="./dataset/" diff --git a/dataset/readme.txt b/dataset/readme.txt new file mode 100755 index 00000000..336d5ad7 --- /dev/null +++ b/dataset/readme.txt @@ -0,0 +1 @@ +Download the preprocessed datasets from [here](https://drive.google.com/drive/folders/112GPdRC9IEcwcJRyrLJeYw9_YV4wLdKC?usp=sharing) and place them in this folder. \ No newline at end of file diff --git a/demo/vis_in_the_wild.py b/demo/vis_in_the_wild.py index 02e572c0..9ca6f1ee 100755 --- a/demo/vis_in_the_wild.py +++ b/demo/vis_in_the_wild.py @@ -19,12 +19,12 @@ sys.path.append(os.getcwd()) # Auto-download checkpoint files if missing -from fmpose.lib.checkpoint.download_checkpoints import ensure_checkpoints +from fmpose3d.lib.checkpoint.download_checkpoints import ensure_checkpoints ensure_checkpoints() -from fmpose.lib.preprocess import h36m_coco_format, revise_kpts -from fmpose.lib.hrnet.gen_kpts import gen_video_kpts as hrnet_pose -from fmpose.common.arguments import opts as parse_args +from fmpose3d.lib.preprocess import h36m_coco_format, revise_kpts +from fmpose3d.lib.hrnet.gen_kpts import gen_video_kpts as hrnet_pose +from fmpose3d.common.arguments import opts as parse_args args = parse_args().parse() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu @@ -39,7 +39,7 @@ spec.loader.exec_module(module) CFM = getattr(module, 'Model') -from fmpose.common.camera import * +from fmpose3d.common.camera import * import matplotlib import matplotlib.pyplot as plt diff --git a/fmpose/animals/__init__.py b/fmpose/animals/__init__.py deleted file mode 100644 index fd963384..00000000 --- a/fmpose/animals/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -""" -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 -""" - -""" -Animal-specific components for FMPose. -""" - -__all__ = [ - "common", -] - diff --git a/fmpose/__init__.py b/fmpose3d/__init__.py similarity index 87% rename from fmpose/__init__.py rename to fmpose3d/__init__.py index 5b3c2116..8a9a4716 100644 --- a/fmpose/__init__.py +++ b/fmpose3d/__init__.py @@ -7,15 +7,15 @@ Licensed under Apache 2.0 """ -__version__ = "0.0.5" +__version__ = "0.0.7" __author__ = "Ti Wang, Xiaohang Yu, Mackenzie Weygandt Mathis" -__license__ = "MIT" +__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_weighted_by_2D_error, + aggregation_RPEA_joint_level, ) # Import 2D pose detection utilities @@ -27,7 +27,7 @@ # Aggregation methods "average_aggregation", "aggregation_select_single_best_hypothesis_by_2D_error", - "aggregation_RPEA_weighted_by_2D_error", + "aggregation_RPEA_joint_level", # 2D pose detection "gen_video_kpts", "h36m_coco_format", diff --git a/fmpose/aggregation_methods.py b/fmpose3d/aggregation_methods.py similarity index 98% rename from fmpose/aggregation_methods.py rename to fmpose3d/aggregation_methods.py index ee2b62d5..38d022f9 100644 --- a/fmpose/aggregation_methods.py +++ b/fmpose3d/aggregation_methods.py @@ -8,7 +8,7 @@ """ import torch -from fmpose.common.utils import project_to_2d +from fmpose3d.common.utils import project_to_2d def average_aggregation(list_hypothesis): return torch.mean(torch.stack(list_hypothesis), dim=0) @@ -96,7 +96,7 @@ def aggregation_select_single_best_hypothesis_by_2D_error(args, return agg -def aggregation_RPEA_weighted_by_2D_error( +def aggregation_RPEA_joint_level( args, list_hypothesis, batch_cam, input_2D, gt_3D, topk=3 ): """ diff --git a/fmpose3d/animals/__init__.py b/fmpose3d/animals/__init__.py new file mode 100644 index 00000000..21ceac08 --- /dev/null +++ b/fmpose3d/animals/__init__.py @@ -0,0 +1,8 @@ +""" +Animal-specific components for FMPose3D. +""" + +__all__ = [ + "common", +] + diff --git a/fmpose/animals/common/__init__.py b/fmpose3d/animals/common/__init__.py similarity index 100% rename from fmpose/animals/common/__init__.py rename to fmpose3d/animals/common/__init__.py diff --git a/fmpose/animals/common/animal3d_dataset.py b/fmpose3d/animals/common/animal3d_dataset.py similarity index 100% rename from fmpose/animals/common/animal3d_dataset.py rename to fmpose3d/animals/common/animal3d_dataset.py diff --git a/fmpose/animals/common/animal_visualization.py b/fmpose3d/animals/common/animal_visualization.py similarity index 100% rename from fmpose/animals/common/animal_visualization.py rename to fmpose3d/animals/common/animal_visualization.py diff --git a/fmpose/animals/common/arber_dataset.py b/fmpose3d/animals/common/arber_dataset.py similarity index 100% rename from fmpose/animals/common/arber_dataset.py rename to fmpose3d/animals/common/arber_dataset.py diff --git a/fmpose/animals/common/arguments.py b/fmpose3d/animals/common/arguments.py similarity index 100% rename from fmpose/animals/common/arguments.py rename to fmpose3d/animals/common/arguments.py diff --git a/fmpose/animals/common/camera.py b/fmpose3d/animals/common/camera.py similarity index 100% rename from fmpose/animals/common/camera.py rename to fmpose3d/animals/common/camera.py diff --git a/fmpose/animals/common/graph_utils.py b/fmpose3d/animals/common/graph_utils.py similarity index 100% rename from fmpose/animals/common/graph_utils.py rename to fmpose3d/animals/common/graph_utils.py diff --git a/fmpose/animals/common/lifter3d.py b/fmpose3d/animals/common/lifter3d.py similarity index 100% rename from fmpose/animals/common/lifter3d.py rename to fmpose3d/animals/common/lifter3d.py diff --git a/fmpose/animals/common/mocap_dataset.py b/fmpose3d/animals/common/mocap_dataset.py similarity index 100% rename from fmpose/animals/common/mocap_dataset.py rename to fmpose3d/animals/common/mocap_dataset.py diff --git a/fmpose/animals/common/skeleton.py b/fmpose3d/animals/common/skeleton.py similarity index 100% rename from fmpose/animals/common/skeleton.py rename to fmpose3d/animals/common/skeleton.py diff --git a/fmpose/animals/common/utils.py b/fmpose3d/animals/common/utils.py similarity index 100% rename from fmpose/animals/common/utils.py rename to fmpose3d/animals/common/utils.py diff --git a/fmpose/animals/models/__init__.py b/fmpose3d/animals/models/__init__.py similarity index 100% rename from fmpose/animals/models/__init__.py rename to fmpose3d/animals/models/__init__.py diff --git a/fmpose/animals/models/graph_frames.py b/fmpose3d/animals/models/graph_frames.py similarity index 100% rename from fmpose/animals/models/graph_frames.py rename to fmpose3d/animals/models/graph_frames.py diff --git a/fmpose/animals/models/model_animal3d.py b/fmpose3d/animals/models/model_animal3d.py similarity index 97% rename from fmpose/animals/models/model_animal3d.py rename to fmpose3d/animals/models/model_animal3d.py index 2fbc1750..273b1ddc 100644 --- a/fmpose/animals/models/model_animal3d.py +++ b/fmpose3d/animals/models/model_animal3d.py @@ -15,7 +15,7 @@ from einops import rearrange from timm.models.layers import DropPath -from fmpose.animals.models.graph_frames import Graph +from fmpose3d.animals.models.graph_frames import Graph class TimeEmbedding(nn.Module): def __init__(self, dim: int, hidden_dim: int = 64): @@ -148,7 +148,7 @@ def forward(self, x): x = res2 + self.drop_path(x) return x -class FMPose(nn.Module): +class FMPose3D(nn.Module): def __init__(self, depth, embed_dim, channels_dim, tokens_dim, adj, drop_rate=0.10, length=27): super().__init__() drop_path_rate = 0.2 @@ -219,7 +219,7 @@ def __init__(self, args): self.encoder_pose_2d = encoder(2, args.channel//2, args.channel//2-self.t_embed_dim//2) self.encoder_y_t = encoder(3, args.channel//2, args.channel//2-self.t_embed_dim//2) - self.FMPose = FMPose(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256 + self.FMPose3D = FMPose3D(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256 self.pred_mu = decoder(args.channel, args.channel//2, 3) def forward(self, pose_2d, y_t, t): @@ -238,7 +238,7 @@ def forward(self, pose_2d, y_t, t): in_emb = rearrange(in_emb, 'b f j c -> (b f) j c').contiguous() # (B*F,J,in) # encoder -> model -> regression head - h = self.FMPose(in_emb) + h = self.FMPose3D(in_emb) v = self.pred_mu(h) # (B*F,J,3) v = rearrange(v, '(b f) j c -> b f j c', b=b, f=f).contiguous() # (B,F,J,3) diff --git a/fmpose/common/__init__.py b/fmpose3d/common/__init__.py similarity index 100% rename from fmpose/common/__init__.py rename to fmpose3d/common/__init__.py diff --git a/fmpose/common/arguments.py b/fmpose3d/common/arguments.py similarity index 100% rename from fmpose/common/arguments.py rename to fmpose3d/common/arguments.py diff --git a/fmpose/common/camera.py b/fmpose3d/common/camera.py similarity index 100% rename from fmpose/common/camera.py rename to fmpose3d/common/camera.py diff --git a/fmpose/common/generator.py b/fmpose3d/common/generator.py similarity index 100% rename from fmpose/common/generator.py rename to fmpose3d/common/generator.py diff --git a/fmpose/common/graph_utils.py b/fmpose3d/common/graph_utils.py similarity index 100% rename from fmpose/common/graph_utils.py rename to fmpose3d/common/graph_utils.py diff --git a/fmpose/common/h36m_dataset.py b/fmpose3d/common/h36m_dataset.py similarity index 98% rename from fmpose/common/h36m_dataset.py rename to fmpose3d/common/h36m_dataset.py index 11e07395..061810e9 100755 --- a/fmpose/common/h36m_dataset.py +++ b/fmpose3d/common/h36m_dataset.py @@ -11,9 +11,9 @@ import numpy as np -from fmpose.common.camera import normalize_screen_coordinates -from fmpose.common.mocap_dataset import MocapDataset -from fmpose.common.skeleton import Skeleton +from fmpose3d.common.camera import normalize_screen_coordinates +from fmpose3d.common.mocap_dataset import MocapDataset +from fmpose3d.common.skeleton import Skeleton h36m_skeleton = Skeleton( parents=[ diff --git a/fmpose/common/load_data_hm36.py b/fmpose3d/common/load_data_hm36.py similarity index 98% rename from fmpose/common/load_data_hm36.py rename to fmpose3d/common/load_data_hm36.py index be38b595..bf411dfb 100755 --- a/fmpose/common/load_data_hm36.py +++ b/fmpose3d/common/load_data_hm36.py @@ -10,9 +10,9 @@ import numpy as np import torch.utils.data as data -from fmpose.common.camera import normalize_screen_coordinates, world_to_camera -from fmpose.common.generator import ChunkedGenerator -from fmpose.common.utils import deterministic_random +from fmpose3d.common.camera import normalize_screen_coordinates, world_to_camera +from fmpose3d.common.generator import ChunkedGenerator +from fmpose3d.common.utils import deterministic_random class Fusion(data.Dataset): diff --git a/fmpose/common/mocap_dataset.py b/fmpose3d/common/mocap_dataset.py similarity index 100% rename from fmpose/common/mocap_dataset.py rename to fmpose3d/common/mocap_dataset.py diff --git a/fmpose/common/skeleton.py b/fmpose3d/common/skeleton.py similarity index 100% rename from fmpose/common/skeleton.py rename to fmpose3d/common/skeleton.py diff --git a/fmpose/common/utils.py b/fmpose3d/common/utils.py similarity index 100% rename from fmpose/common/utils.py rename to fmpose3d/common/utils.py diff --git a/fmpose/lib/__init__.py b/fmpose3d/lib/__init__.py similarity index 100% rename from fmpose/lib/__init__.py rename to fmpose3d/lib/__init__.py diff --git a/fmpose/lib/checkpoint/README.md b/fmpose3d/lib/checkpoint/README.md similarity index 100% rename from fmpose/lib/checkpoint/README.md rename to fmpose3d/lib/checkpoint/README.md diff --git a/fmpose/lib/checkpoint/__init__.py b/fmpose3d/lib/checkpoint/__init__.py similarity index 100% rename from fmpose/lib/checkpoint/__init__.py rename to fmpose3d/lib/checkpoint/__init__.py diff --git a/fmpose/lib/checkpoint/download_checkpoints.py b/fmpose3d/lib/checkpoint/download_checkpoints.py similarity index 100% rename from fmpose/lib/checkpoint/download_checkpoints.py rename to fmpose3d/lib/checkpoint/download_checkpoints.py diff --git a/fmpose/lib/hrnet/__init__.py b/fmpose3d/lib/hrnet/__init__.py similarity index 100% rename from fmpose/lib/hrnet/__init__.py rename to fmpose3d/lib/hrnet/__init__.py diff --git a/fmpose/lib/hrnet/experiments/w48_384x288_adam_lr1e-3.yaml b/fmpose3d/lib/hrnet/experiments/w48_384x288_adam_lr1e-3.yaml similarity index 100% rename from fmpose/lib/hrnet/experiments/w48_384x288_adam_lr1e-3.yaml rename to fmpose3d/lib/hrnet/experiments/w48_384x288_adam_lr1e-3.yaml diff --git a/fmpose/lib/hrnet/gen_kpts.py b/fmpose3d/lib/hrnet/gen_kpts.py similarity index 91% rename from fmpose/lib/hrnet/gen_kpts.py rename to fmpose3d/lib/hrnet/gen_kpts.py index 1ae3f79c..a75e7046 100755 --- a/fmpose/lib/hrnet/gen_kpts.py +++ b/fmpose3d/lib/hrnet/gen_kpts.py @@ -24,22 +24,22 @@ import cv2 import copy -from fmpose.lib.hrnet.lib.utils.utilitys import plot_keypoint, PreProcess, write, load_json -from fmpose.lib.hrnet.lib.config import cfg, update_config -from fmpose.lib.hrnet.lib.utils.transforms import * -from fmpose.lib.hrnet.lib.utils.inference import get_final_preds -from fmpose.lib.hrnet.lib.models import pose_hrnet +from fmpose3d.lib.hrnet.lib.utils.utilitys import plot_keypoint, PreProcess, write, load_json +from fmpose3d.lib.hrnet.lib.config import cfg, update_config +from fmpose3d.lib.hrnet.lib.utils.transforms import * +from fmpose3d.lib.hrnet.lib.utils.inference import get_final_preds +from fmpose3d.lib.hrnet.lib.models import pose_hrnet cfg_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'experiments') + '/' # Auto-download checkpoints if missing and get checkpoint paths -from fmpose.lib.checkpoint.download_checkpoints import ensure_checkpoints, get_checkpoint_path +from fmpose3d.lib.checkpoint.download_checkpoints import ensure_checkpoints, get_checkpoint_path ensure_checkpoints() # Loading human detector model -from fmpose.lib.yolov3.human_detector import load_model as yolo_model -from fmpose.lib.yolov3.human_detector import yolo_human_det as yolo_det -from fmpose.lib.sort.sort import Sort +from fmpose3d.lib.yolov3.human_detector import load_model as yolo_model +from fmpose3d.lib.yolov3.human_detector import yolo_human_det as yolo_det +from fmpose3d.lib.sort.sort import Sort def parse_args(): parser = argparse.ArgumentParser(description='Train keypoints network') diff --git a/fmpose/lib/hrnet/lib/__init__.py b/fmpose3d/lib/hrnet/lib/__init__.py similarity index 100% rename from fmpose/lib/hrnet/lib/__init__.py rename to fmpose3d/lib/hrnet/lib/__init__.py diff --git a/fmpose/lib/hrnet/lib/config/__init__.py b/fmpose3d/lib/hrnet/lib/config/__init__.py similarity index 100% rename from fmpose/lib/hrnet/lib/config/__init__.py rename to fmpose3d/lib/hrnet/lib/config/__init__.py diff --git a/fmpose/lib/hrnet/lib/config/default.py b/fmpose3d/lib/hrnet/lib/config/default.py similarity index 100% rename from fmpose/lib/hrnet/lib/config/default.py rename to fmpose3d/lib/hrnet/lib/config/default.py diff --git a/fmpose/lib/hrnet/lib/config/models.py b/fmpose3d/lib/hrnet/lib/config/models.py similarity index 100% rename from fmpose/lib/hrnet/lib/config/models.py rename to fmpose3d/lib/hrnet/lib/config/models.py diff --git a/fmpose/lib/hrnet/lib/models/__init__.py b/fmpose3d/lib/hrnet/lib/models/__init__.py similarity index 100% rename from fmpose/lib/hrnet/lib/models/__init__.py rename to fmpose3d/lib/hrnet/lib/models/__init__.py diff --git a/fmpose/lib/hrnet/lib/models/pose_hrnet.py b/fmpose3d/lib/hrnet/lib/models/pose_hrnet.py similarity index 100% rename from fmpose/lib/hrnet/lib/models/pose_hrnet.py rename to fmpose3d/lib/hrnet/lib/models/pose_hrnet.py diff --git a/fmpose/lib/hrnet/lib/utils/__init__.py b/fmpose3d/lib/hrnet/lib/utils/__init__.py similarity index 100% rename from fmpose/lib/hrnet/lib/utils/__init__.py rename to fmpose3d/lib/hrnet/lib/utils/__init__.py diff --git a/fmpose/lib/hrnet/lib/utils/coco_h36m.py b/fmpose3d/lib/hrnet/lib/utils/coco_h36m.py similarity index 100% rename from fmpose/lib/hrnet/lib/utils/coco_h36m.py rename to fmpose3d/lib/hrnet/lib/utils/coco_h36m.py diff --git a/fmpose/lib/hrnet/lib/utils/inference.py b/fmpose3d/lib/hrnet/lib/utils/inference.py similarity index 100% rename from fmpose/lib/hrnet/lib/utils/inference.py rename to fmpose3d/lib/hrnet/lib/utils/inference.py diff --git a/fmpose/lib/hrnet/lib/utils/transforms.py b/fmpose3d/lib/hrnet/lib/utils/transforms.py similarity index 100% rename from fmpose/lib/hrnet/lib/utils/transforms.py rename to fmpose3d/lib/hrnet/lib/utils/transforms.py diff --git a/fmpose/lib/hrnet/lib/utils/utilitys.py b/fmpose3d/lib/hrnet/lib/utils/utilitys.py similarity index 98% rename from fmpose/lib/hrnet/lib/utils/utilitys.py rename to fmpose3d/lib/hrnet/lib/utils/utilitys.py index 5f4271f9..fde7b3de 100755 --- a/fmpose/lib/hrnet/lib/utils/utilitys.py +++ b/fmpose3d/lib/hrnet/lib/utils/utilitys.py @@ -12,9 +12,9 @@ import torch import json import torchvision.transforms as transforms -from fmpose.lib.hrnet.lib.utils.transforms import * +from fmpose3d.lib.hrnet.lib.utils.transforms import * -from fmpose.lib.hrnet.lib.utils.coco_h36m import coco_h36m +from fmpose3d.lib.hrnet.lib.utils.coco_h36m import coco_h36m import numpy as np joint_pairs = [[0, 1], [1, 3], [0, 2], [2, 4], diff --git a/fmpose/lib/preprocess.py b/fmpose3d/lib/preprocess.py similarity index 100% rename from fmpose/lib/preprocess.py rename to fmpose3d/lib/preprocess.py diff --git a/fmpose/lib/sort/__init__.py b/fmpose3d/lib/sort/__init__.py similarity index 100% rename from fmpose/lib/sort/__init__.py rename to fmpose3d/lib/sort/__init__.py diff --git a/fmpose/lib/sort/sort.py b/fmpose3d/lib/sort/sort.py similarity index 100% rename from fmpose/lib/sort/sort.py rename to fmpose3d/lib/sort/sort.py diff --git a/fmpose/lib/yolov3/__init__.py b/fmpose3d/lib/yolov3/__init__.py similarity index 100% rename from fmpose/lib/yolov3/__init__.py rename to fmpose3d/lib/yolov3/__init__.py diff --git a/fmpose/lib/yolov3/bbox.py b/fmpose3d/lib/yolov3/bbox.py similarity index 100% rename from fmpose/lib/yolov3/bbox.py rename to fmpose3d/lib/yolov3/bbox.py diff --git a/fmpose/lib/yolov3/cfg/tiny-yolo-voc.cfg b/fmpose3d/lib/yolov3/cfg/tiny-yolo-voc.cfg similarity index 100% rename from fmpose/lib/yolov3/cfg/tiny-yolo-voc.cfg rename to fmpose3d/lib/yolov3/cfg/tiny-yolo-voc.cfg diff --git a/fmpose/lib/yolov3/cfg/yolo-voc.cfg b/fmpose3d/lib/yolov3/cfg/yolo-voc.cfg similarity index 100% rename from fmpose/lib/yolov3/cfg/yolo-voc.cfg rename to fmpose3d/lib/yolov3/cfg/yolo-voc.cfg diff --git a/fmpose/lib/yolov3/cfg/yolo.cfg b/fmpose3d/lib/yolov3/cfg/yolo.cfg similarity index 100% rename from fmpose/lib/yolov3/cfg/yolo.cfg rename to fmpose3d/lib/yolov3/cfg/yolo.cfg diff --git a/fmpose/lib/yolov3/cfg/yolov3.cfg b/fmpose3d/lib/yolov3/cfg/yolov3.cfg similarity index 100% rename from fmpose/lib/yolov3/cfg/yolov3.cfg rename to fmpose3d/lib/yolov3/cfg/yolov3.cfg diff --git a/fmpose/lib/yolov3/darknet.py b/fmpose3d/lib/yolov3/darknet.py similarity index 99% rename from fmpose/lib/yolov3/darknet.py rename to fmpose3d/lib/yolov3/darknet.py index 941316be..a5e6e7d6 100755 --- a/fmpose/lib/yolov3/darknet.py +++ b/fmpose3d/lib/yolov3/darknet.py @@ -17,8 +17,8 @@ import os import sys -from fmpose.lib.yolov3.util import convert2cpu as cpu -from fmpose.lib.yolov3.util import predict_transform +from fmpose3d.lib.yolov3.util import convert2cpu as cpu +from fmpose3d.lib.yolov3.util import predict_transform class test_net(nn.Module): diff --git a/fmpose/lib/yolov3/data/coco.names b/fmpose3d/lib/yolov3/data/coco.names similarity index 100% rename from fmpose/lib/yolov3/data/coco.names rename to fmpose3d/lib/yolov3/data/coco.names diff --git a/fmpose/lib/yolov3/data/pallete b/fmpose3d/lib/yolov3/data/pallete similarity index 100% rename from fmpose/lib/yolov3/data/pallete rename to fmpose3d/lib/yolov3/data/pallete diff --git a/fmpose/lib/yolov3/data/voc.names b/fmpose3d/lib/yolov3/data/voc.names similarity index 100% rename from fmpose/lib/yolov3/data/voc.names rename to fmpose3d/lib/yolov3/data/voc.names diff --git a/fmpose/lib/yolov3/human_detector.py b/fmpose3d/lib/yolov3/human_detector.py similarity index 96% rename from fmpose/lib/yolov3/human_detector.py rename to fmpose3d/lib/yolov3/human_detector.py index c3eecce0..5f9cdab2 100755 --- a/fmpose/lib/yolov3/human_detector.py +++ b/fmpose3d/lib/yolov3/human_detector.py @@ -18,10 +18,10 @@ import pickle as pkl import argparse -from fmpose.lib.yolov3.util import * -from fmpose.lib.yolov3.darknet import Darknet -from fmpose.lib.yolov3 import preprocess -from fmpose.lib.checkpoint.download_checkpoints import get_checkpoint_path +from fmpose3d.lib.yolov3.util import * +from fmpose3d.lib.yolov3.darknet import Darknet +from fmpose3d.lib.yolov3 import preprocess +from fmpose3d.lib.checkpoint.download_checkpoints import get_checkpoint_path cur_dir = os.path.dirname(os.path.realpath(__file__)) project_root = os.path.join(cur_dir, '../../../') diff --git a/fmpose/lib/yolov3/preprocess.py b/fmpose3d/lib/yolov3/preprocess.py similarity index 100% rename from fmpose/lib/yolov3/preprocess.py rename to fmpose3d/lib/yolov3/preprocess.py diff --git a/fmpose/lib/yolov3/util.py b/fmpose3d/lib/yolov3/util.py similarity index 98% rename from fmpose/lib/yolov3/util.py rename to fmpose3d/lib/yolov3/util.py index 79808ba9..1db0d509 100755 --- a/fmpose/lib/yolov3/util.py +++ b/fmpose3d/lib/yolov3/util.py @@ -13,8 +13,8 @@ import numpy as np import cv2 import os.path as osp -from fmpose.lib.yolov3.bbox import bbox_iou -from fmpose.lib.checkpoint.download_checkpoints import get_checkpoint_dir +from fmpose3d.lib.yolov3.bbox import bbox_iou +from fmpose3d.lib.checkpoint.download_checkpoints import get_checkpoint_dir def get_path(cur_file): diff --git a/fmpose/models/__init__.py b/fmpose3d/models/__init__.py similarity index 95% rename from fmpose/models/__init__.py rename to fmpose3d/models/__init__.py index bd33b79a..5b94df4d 100644 --- a/fmpose/models/__init__.py +++ b/fmpose3d/models/__init__.py @@ -8,7 +8,7 @@ """ """ -FMPose models. +FMPose3D models. """ from .graph_frames import Graph diff --git a/fmpose/models/graph_frames.py b/fmpose3d/models/graph_frames.py similarity index 100% rename from fmpose/models/graph_frames.py rename to fmpose3d/models/graph_frames.py diff --git a/fmpose/models/model_GAMLP.py b/fmpose3d/models/model_GAMLP.py similarity index 97% rename from fmpose/models/model_GAMLP.py rename to fmpose3d/models/model_GAMLP.py index 9a016d4a..c0c6b46e 100644 --- a/fmpose/models/model_GAMLP.py +++ b/fmpose3d/models/model_GAMLP.py @@ -13,7 +13,7 @@ import torch.nn as nn import math from einops import rearrange -from fmpose.models.graph_frames import Graph +from fmpose3d.models.graph_frames import Graph from functools import partial from einops import rearrange, repeat from timm.models.layers import DropPath @@ -154,7 +154,7 @@ def forward(self, x): x = res2 + self.drop_path(x) return x -class FMPose(nn.Module): +class FMPose3D(nn.Module): def __init__(self, depth, embed_dim, channels_dim, tokens_dim, adj, drop_rate=0.10, length=27): super().__init__() @@ -228,7 +228,7 @@ def __init__(self, args): self.encoder_pose_2d = encoder(2, args.channel//2, args.channel//2-self.t_embed_dim//2) self.encoder_y_t = encoder(3, args.channel//2, args.channel//2-self.t_embed_dim//2) - self.FMPose = FMPose(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256 + self.FMPose3D = FMPose3D(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256 self.pred_mu = decoder(args.channel, args.channel//2, 3) def forward(self, pose_2d, y_t, t): @@ -250,7 +250,7 @@ def forward(self, pose_2d, y_t, t): in_emb = rearrange(in_emb, 'b f j c -> (b f) j c').contiguous() # (B*F,J,in) # encoder -> model -> regression head - h = self.FMPose(in_emb) + h = self.FMPose3D(in_emb) v = self.pred_mu(h) # (B*F,J,3) v = rearrange(v, '(b f) j c -> b f j c', b=b, f=f).contiguous() # (B,F,J,3) diff --git a/pyproject.toml b/pyproject.toml index 5ede62e0..bd3e8a0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,24 +3,22 @@ requires = ["setuptools>=61.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "fmpose" +name = "fmpose3d" readme = "README.md" requires-python = ">=3.8" dynamic = ["version"] -license = {text = "MIT"} +license = {text = "Apache 2.0"} authors = [ {name = "Ti Wang", email = "ti.wang@epfl.ch"}, {name = "Xiaohang Yu"}, - {name = "Mackenzie Weygandt Mathis"}, + {name = "Mackenzie Weygandt Mathis", email = "mackenzie.mathis@epfl.ch"}, ] keywords = ["pose estimation", "3D pose", "flow matching", "computer vision"] classifiers = [ "Development Status :: 4 - Beta", "Intended Audience :: Science/Research", - "License :: OSI Approved :: MIT License", + "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering :: Artificial Intelligence", ] @@ -30,7 +28,7 @@ dependencies = [ "torchvision==0.19.1", "timm>=1.0.0", "einops==0.4.0", - "numpy==2.1.1", + "numpy>=1.18.5,<2.0", "tqdm>=4.60.0", "scipy>=1.7.0", "yacs>=0.1.8", @@ -38,6 +36,8 @@ dependencies = [ "numba>=0.56.0", "scikit-image>=0.19.0", "filterpy>=1.4.5", + "pandas>=1.0.1", + "deeplabcut==3.0.0rc13", ] [project.optional-dependencies] @@ -56,7 +56,7 @@ include = ["fmpose*"] "*" = ["*.yaml", "*.cfg", "*.names"] [tool.setuptools.dynamic] -version = {attr = "fmpose.__version__"} +version = {attr = "fmpose3d.__version__"} [tool.black] line-length = 100 diff --git a/scripts/FMPose_main.py b/scripts/FMPose3D_main.py similarity index 98% rename from scripts/FMPose_main.py rename to scripts/FMPose3D_main.py index b0df3376..6cd97b1d 100644 --- a/scripts/FMPose_main.py +++ b/scripts/FMPose3D_main.py @@ -18,10 +18,10 @@ import torch.optim as optim from tqdm import tqdm -from fmpose.common import opts, Human36mDataset, Fusion -from fmpose.common.utils import * +from fmpose3d.common import opts, Human36mDataset, Fusion +from fmpose3d.common.utils import * -from fmpose.aggregation_methods import aggregation_RPEA_weighted_by_2D_error +from fmpose3d.aggregation_methods import aggregation_RPEA_joint_level args = opts().parse() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu @@ -42,7 +42,7 @@ CFM = getattr(module, "Model") else: # Load model from installed fmpose package - from fmpose.models import Model as CFM + from fmpose3d.models import Model as CFM def test_multi_hypothesis( @@ -130,7 +130,7 @@ def euler_sample(x2d, y_local, steps): output_3D_s[:, :, 0, :] = 0 list_hypothesis.append(output_3D_s) - output_3D_s = aggregation_RPEA_weighted_by_2D_error(args, + output_3D_s = aggregation_RPEA_joint_level(args, list_hypothesis, batch_cam, input_2D_nonflip, gt_3D, args.topk ) diff --git a/scripts/FMPose_test.sh b/scripts/FMPose3D_test.sh similarity index 91% rename from scripts/FMPose_test.sh rename to scripts/FMPose3D_test.sh index 985f9a82..de7869ea 100755 --- a/scripts/FMPose_test.sh +++ b/scripts/FMPose3D_test.sh @@ -1,7 +1,7 @@ #inference layers=5 batch_size=1024 -sh_file='scripts/FMPose_test.sh' +sh_file='scripts/FMPose3D_test.sh' weight_softmax_tau=1.0 num_hypothesis_list=1 eval_multi_steps=3 @@ -15,7 +15,7 @@ model_path='pre_trained_models/fmpose_detected2d/model_GAMLP.py' saved_model_path='pre_trained_models/fmpose_detected2d/FMpose_36_4972_best.pth' #Test CFM -python3 scripts/FMPose_main.py \ +python3 scripts/FMPose3D_main.py \ --reload \ --topk ${topk} \ --exp_temp ${exp_temp} \ diff --git a/scripts/FMPose_train.sh b/scripts/FMPose3D_train.sh similarity index 85% rename from scripts/FMPose_train.sh rename to scripts/FMPose3D_train.sh index 613e6040..cb3e0289 100755 --- a/scripts/FMPose_train.sh +++ b/scripts/FMPose3D_train.sh @@ -1,4 +1,4 @@ -#Train FMPose +#Train FMPose3D layers=5 lr=1e-3 decay=0.98 @@ -13,10 +13,10 @@ frames=1 channel_dim=512 model_path="" # when the path is empty, the model will be loaded from the installed fmpose package # model_path='./fmpose/models/model_GAMLP.py' # when the path is not empty, the model will be loaded from the local file path -sh_file='scripts/FMPose_train.sh' -folder_name=FMPose_Publish_layers${layers}_$(date +%Y%m%d_%H%M%S) +sh_file='scripts/FMPose3D_train.sh' +folder_name=FMPose3D_Publish_layers${layers}_$(date +%Y%m%d_%H%M%S) -python3 scripts/FMPose_main.py \ +python3 scripts/FMPose3D_main.py \ --train \ --dataset h36m \ --frames ${frames} \ diff --git a/tests/test_demo_human.py b/tests/test_demo_human.py index c0728189..ae8b71af 100644 --- a/tests/test_demo_human.py +++ b/tests/test_demo_human.py @@ -34,8 +34,8 @@ def test_output_dir(tmp_path): def test_2d_pose_estimation(test_image_path, test_output_dir): """Test that 2D pose estimation runs and produces output.""" # Import here to avoid import issues at collection time - from fmpose.lib.hrnet.gen_kpts import gen_video_kpts as hrnet_pose - from fmpose.lib.preprocess import h36m_coco_format, revise_kpts + from fmpose3d.lib.hrnet.gen_kpts import gen_video_kpts as hrnet_pose + from fmpose3d.lib.preprocess import h36m_coco_format, revise_kpts # Run 2D pose estimation keypoints, scores = hrnet_pose(test_image_path, det_dim=416, num_peroson=1, gen_output=True, type='image') @@ -61,9 +61,9 @@ def test_2d_pose_estimation(test_image_path, test_output_dir): def test_demo_pipeline_runs(test_image_path): """Test that the full demo pipeline can be imported and key components work.""" # Test imports - from fmpose.lib.hrnet.gen_kpts import gen_video_kpts - from fmpose.lib.preprocess import h36m_coco_format, revise_kpts - from fmpose.models import Model + from fmpose3d.lib.hrnet.gen_kpts import gen_video_kpts + from fmpose3d.lib.preprocess import h36m_coco_format, revise_kpts + from fmpose3d.models import Model assert gen_video_kpts is not None assert h36m_coco_format is not None diff --git a/tests/test_model.py b/tests/test_model.py index 7e8edbb5..09729d38 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -9,7 +9,7 @@ import pytest import torch -from fmpose.models import Model +from fmpose3d.models import Model class Args: """Mock args for model configuration.""" diff --git a/tests/test_training_pipeline.py b/tests/test_training_pipeline.py index d0a79289..f9f59c64 100644 --- a/tests/test_training_pipeline.py +++ b/tests/test_training_pipeline.py @@ -10,7 +10,7 @@ import pytest import torch import torch.optim as optim -from fmpose.models import Model +from fmpose3d.models import Model class Args: @@ -46,7 +46,7 @@ def optimizer(model): def sample_batch(device): """Create a sample batch for training. - Shapes based on FMPose_main.py train function: + Shapes based on FMPose3D_main.py train function: - input_2D: (B, F, J, 2) - 2D pose input - gt_3D: (B, F, J, 3) - 3D pose ground truth """