Add Docker support and update test/demo scripts to use automatic Hugging Face downloads

Dockerfile

@@ -0,0 +1,64 @@
+FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime
+
+ARG USERNAME=fmpose3d
+ARG USER_UID=1000
+ARG USER_GID=1000
+RUN groupadd ${USERNAME} --gid ${USER_GID}
+RUN useradd -m -s /bin/bash -g ${USERNAME} -u ${USER_UID} ${USERNAME}
+RUN mkdir /app /logs /data
+
+ENV DEBIAN_FRONTEND=noninteractive
+SHELL ["/bin/bash", "-c"]
+RUN apt-get update -y && apt-get install -yy --no-install-recommends \
+    vim zsh tmux wget curl htop git sudo ssh git-lfs \
+    python3 python3-pip \
+    libgl1-mesa-glx libglib2.0-0 \
+    ffmpeg \
+    && apt-get -y autoclean \
+    && apt-get -y autoremove \
+    && rm -rf /var/lib/apt/lists/* \
+    && apt-get clean
+
+ENV PATH="/opt/conda/bin:${PATH}"
+
+# Initialize conda for root just in case and fix symlinks
+RUN ln -fs /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh
+
+# --- Install FMPose3D from PyPI (as documented in README) ---
+RUN python -m pip install --no-cache-dir --upgrade pip && \
+    python -m pip install --no-cache-dir "fmpose3d[animals,viz]" gdown
+
+# Allow non-root user to download DLC model weights at runtime
+RUN mkdir -p /opt/conda/lib/python3.11/site-packages/deeplabcut/modelzoo/checkpoints && \
+    chown -R ${USERNAME}:${USERNAME} /opt/conda/lib/python3.11/site-packages/deeplabcut/modelzoo
+
+# Set your user as owner of the home directory before switching
+RUN chown -R ${USERNAME}:${USERNAME} /home/${USERNAME}
+
+ENV NVIDIA_DRIVER_CAPABILITIES=all
+
+# --- SWITCH TO USER ---
+USER ${USERNAME}
+ENV HOME=/home/${USERNAME}
+WORKDIR ${HOME}
+
+# Ensure dotfiles exist
+RUN touch ${HOME}/.bashrc ${HOME}/.zshrc && \
+    chmod 644 ${HOME}/.bashrc ${HOME}/.zshrc
+
+# Install Oh My Zsh and plugins (MUST come before conda init, since OMZ overwrites .zshrc)
+RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" --unattended \
+    && git clone https://github.com/zsh-users/zsh-autosuggestions ~/.oh-my-zsh/custom/plugins/zsh-autosuggestions \
+    && git clone https://github.com/zsh-users/zsh-syntax-highlighting ~/.oh-my-zsh/custom/plugins/zsh-syntax-highlighting \
+    && sed -i 's/^plugins=(.*)$/plugins=(git zsh-autosuggestions zsh-syntax-highlighting)/' ~/.zshrc \
+    && echo "export PATH=\$PATH:${HOME}/.local/bin" >> ~/.zshrc
+
+# Initialize conda AFTER Oh My Zsh (so conda init block is not overwritten)
+RUN /opt/conda/bin/conda init bash && /opt/conda/bin/conda init zsh
+
+# Add conda activation to bashrc and zshrc
+RUN echo "conda activate base" >> ${HOME}/.bashrc && \
+    echo "conda activate base" >> ${HOME}/.zshrc
+
+SHELL ["/bin/zsh", "-c"]
+CMD ["zsh"]

Makefile

@@ -0,0 +1,79 @@
+# [USER: ADJUST PATH]
+PROJECT_NAME := fmpose3d
+project_name_lo := $(shell echo $(PROJECT_NAME) | tr '[:upper:]' '[:lower:]')
+#######
+# RUN #
+#######
+
+# for local server
+IMG_NAME := fmpose3d
+IMG_TAG := v0.1
+DOCKERFILE := Dockerfile
+HOST_UID := $(shell id -u)
+HOST_GID := $(shell id -g)
+BUILD_ARGS := \
+		--build-arg USERNAME=fmpose3d \
+        --build-arg USER_GID=$(HOST_GID) \
+        --build-arg USER_UID=$(HOST_UID)
+
+build:
+	docker build $(BUILD_ARGS) \
+		 -t $(IMG_NAME):$(IMG_TAG) -f $(DOCKERFILE) .
+
+build-clean:
+	docker build --no-cache $(BUILD_ARGS) \
+		 -t $(IMG_NAME):$(IMG_TAG) -f $(DOCKERFILE) .
+
+
+CONTAINER_NAME := fmpose3d_dev1
+# [USER: ADJUST] Mount the project root into the container
+HOST_SRC := $(shell pwd)
+DOCKER_SRC := /fmpose3d
+VOLUMES := \
+	--volume $(HOST_SRC):$(DOCKER_SRC) 
+
+
+run:
+	docker run -it --gpus all --shm-size=64g --name $(CONTAINER_NAME) -w $(DOCKER_SRC) $(VOLUMES) $(IMG_NAME):$(IMG_TAG) 
+
+exec:
+	docker exec -it -w $(DOCKER_SRC) $(CONTAINER_NAME) /bin/zsh
+
+exec_bash:
+	docker exec -it -w $(DOCKER_SRC) $(CONTAINER_NAME) /bin/bash
+
+stop:
+	docker stop $(CONTAINER_NAME)
+
+rm:
+	docker rm $(CONTAINER_NAME)
+
+
+#######################
+# BUILD CORE WITH SRC #
+#######################
+# [USER: ADJUST SRC PATH]
+SRC_PATH="/home/user/project"
+build_production:
+	cp -r $(SRC_PATH) src
+	docker build $(BUILD_ARGS) --build-arg src=src \
+		-t $(IMG_NAME):$(IMG_TAG) -f $(DOCKERFILE) .
+	rm -r src
+
+# Help message
+help:
+	@echo "Available targets:"
+	@echo "  build             - Build Docker image for local server."
+	@echo "  run               - Run a Docker container with GPU."
+	@echo "  exec              - Attach to running container (zsh)."
+	@echo "  exec_bash         - Attach to running container (bash)."
+	@echo "  stop              - Stop the running container."
+	@echo "  rm                - Remove the stopped container."
+	@echo "  build_production  - Build a Docker image with source code."
+	@echo "  help              - Show this help message."
+	@echo ""
+	@echo "Usage:"
+	@echo "  1. make build"
+	@echo "  2. make run"
+	@echo "  3. Inside container: pip install -e '.[animals,viz]'"
+	@echo "  4. Inside container: sh scripts/FMPose3D_train.sh"

demo/vis_in_the_wild.py

@@ -113,7 +113,7 @@ def show3Dpose(vals, ax):
 def get_pose2D(path, output_dir, type):
 
     print('\nGenerating 2D pose...')
-    keypoints, scores = hrnet_pose(path, det_dim=416, num_peroson=1, gen_output=True, type=type)
+    keypoints, scores = hrnet_pose(path, det_dim=416, num_person=1, gen_output=True, type=type)
     keypoints, scores, valid_frames = h36m_coco_format(keypoints, scores)
     re_kpts = revise_kpts(keypoints, scores, valid_frames)
     print('Generating 2D pose successful!')
@@ -278,8 +278,10 @@ def get_pose3D(path, output_dir, type='image'):
 
     # if args.reload:
     model_dict = model['CFM'].state_dict()
-    model_path = args.model_weights_path
-    print(model_path)
+    from fmpose3d.utils.weights import resolve_weights_path
+    model_path = resolve_weights_path(args.model_weights_path, args.model_type)
+
+    print(f"Loading weights from: {model_path}")
     pre_dict = torch.load(model_path, map_location=device, weights_only=True)
     for name, key in model_dict.items():
         model_dict[name] = pre_dict[name]

demo/vis_in_the_wild.sh

@@ -6,7 +6,11 @@ batch_size=1
 sh_file='vis_in_the_wild.sh'
 
 model_type='fmpose3d_humans'
-model_weights_path='../pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
+
+# By default, weights are automatically downloaded from Hugging Face Hub.
+# To use local weights instead, uncomment the line below:
+# model_weights_path='../pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
+model_weights_path=''
 
 target_path='./images/'  # folder containing multiple images
 # target_path='./images/xx.png'  # single image

fmpose3d/inference_api/fmpose3d.py

@@ -34,8 +34,7 @@
 ProgressCallback = Callable[[int, int], None]
 
 
-#: HuggingFace repository hosting the official FMPose3D checkpoints.
-_HF_REPO_ID: str = "deruyter92/fmpose_temp"
+from fmpose3d.utils.weights import HF_REPO_ID as _HF_REPO_ID
 
 # Default camera-to-world rotation quaternion (from the demo script).
 _DEFAULT_CAM_ROTATION = np.array(

fmpose3d/lib/hrnet/gen_kpts.py

@@ -34,7 +34,6 @@
 
 # Auto-download checkpoints if missing and get checkpoint paths
 from fmpose3d.lib.checkpoint.download_checkpoints import ensure_checkpoints, get_checkpoint_path
-ensure_checkpoints()
 
 # Loading human detector model
 from fmpose3d.lib.yolov3.human_detector import load_model as yolo_model
@@ -100,7 +99,7 @@ def model_load(config):
 
     return model
 
-def gen_from_image(args, frame, people_sort, human_model, pose_model, det_dim=416, num_peroson=1, gen_output=False):
+def gen_from_image(args, frame, people_sort, human_model, pose_model, det_dim=416, num_person=1, gen_output=False):
 
     bboxs, scores = yolo_det(frame, human_model, reso=det_dim, confidence=args.thred_score)
     if bboxs is None or not bboxs.any():
@@ -118,7 +117,7 @@ def gen_from_image(args, frame, people_sort, human_model, pose_model, det_dim=41
     if people_track.shape[0] == 1:
         people_track_ = people_track[-1, :-1].reshape(1, 4)
     elif people_track.shape[0] >= 2:
-        people_track_ = people_track[-num_peroson:, :-1].reshape(num_peroson, 4)
+        people_track_ = people_track[-num_person:, :-1].reshape(num_person, 4)
         people_track_ = people_track_[::-1]
     else:
         return [], []
@@ -130,7 +129,7 @@ def gen_from_image(args, frame, people_sort, human_model, pose_model, det_dim=41
 
     with torch.no_grad():
         # bbox is coordinate location
-        inputs, origin_img, center, scale = PreProcess(frame, track_bboxs, cfg, num_peroson)
+        inputs, origin_img, center, scale = PreProcess(frame, track_bboxs, cfg, num_person)
 
         inputs = inputs[:, [2, 1, 0]]
 
@@ -141,8 +140,8 @@ def gen_from_image(args, frame, people_sort, human_model, pose_model, det_dim=41
         # compute coordinate
         preds, maxvals = get_final_preds(cfg, output.clone().cpu().numpy(), np.asarray(center), np.asarray(scale))
 
-    kpts = np.zeros((num_peroson, 17, 2), dtype=np.float32)
-    scores = np.zeros((num_peroson, 17), dtype=np.float32)
+    kpts = np.zeros((num_person, 17, 2), dtype=np.float32)
+    scores = np.zeros((num_person, 17), dtype=np.float32)
     for i, kpt in enumerate(preds):
         kpts[i] = kpt
 
@@ -152,7 +151,10 @@ def gen_from_image(args, frame, people_sort, human_model, pose_model, det_dim=41
     return kpts, scores
 
 
-def gen_video_kpts(path, det_dim=416, num_peroson=1, gen_output=False, type='image'):
+def gen_video_kpts(path, det_dim=416, num_person=1, gen_output=False, type='image'):
+    # Ensure checkpoints are downloaded only when HRNet is actually requested
+    ensure_checkpoints()
+
     # Updating configuration
     args1 = parse_args()
     reset_config(args1)
@@ -166,7 +168,7 @@ def gen_video_kpts(path, det_dim=416, num_peroson=1, gen_output=False, type='ima
     scores_result = []
     if type == "image":
         frame = cv2.imread(path)
-        kpts, scores = gen_from_image(args1, frame, people_sort, human_model, pose_model, det_dim=det_dim, num_peroson=num_peroson, gen_output=gen_output)
+        kpts, scores = gen_from_image(args1, frame, people_sort, human_model, pose_model, det_dim=det_dim, num_person=num_person, gen_output=gen_output)
         kpts_result.append(kpts)
         scores_result.append(scores)
 
@@ -178,7 +180,7 @@ def gen_video_kpts(path, det_dim=416, num_peroson=1, gen_output=False, type='ima
             if not ret:
                 continue
 
-            kpts, scores = gen_from_image(args1, frame, people_sort, human_model, pose_model, det_dim=det_dim, num_peroson=num_peroson, gen_output=gen_output)
+            kpts, scores = gen_from_image(args1, frame, people_sort, human_model, pose_model, det_dim=det_dim, num_person=num_person, gen_output=gen_output)
             kpts_result.append(kpts)
             scores_result.append(scores)
 

fmpose3d/utils/__init__.py

@@ -0,0 +1,8 @@
+"""
+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
+"""

fmpose3d/utils/weights.py

@@ -0,0 +1,49 @@
+"""
+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 helpers for resolving / downloading FMPose3D model weights."""
+
+HF_REPO_ID: str = "deruyter92/fmpose_temp"
+
+
+def resolve_weights_path(model_weights_path: str, model_type: str) -> str:
+    """Return a local weights path, downloading from Hugging Face Hub if needed.
+
+    Parameters
+    ----------
+    model_weights_path : str
+        User-supplied local path.  If falsy the weights are fetched from the
+        Hugging Face Hub automatically.
+    model_type : str
+        Model variant name used to derive the remote filename
+        (e.g. ``"fmpose3d_humans"`` -> ``fmpose3d_humans.pth``).
+
+    Returns
+    -------
+    str
+        Absolute path to the weight file on disk.
+    """
+    if model_weights_path:
+        return model_weights_path
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        raise ImportError(
+            "huggingface_hub is required to download model weights. "
+            "Install it with:  pip install huggingface_hub\n"
+            "Or download the weights manually and pass the local path."
+        ) from None
+
+    filename = f"{model_type}.pth"
+    print(
+        f"No local weights path specified. "
+        f"Downloading '{filename}' from Hugging Face ({HF_REPO_ID})..."
+    )
+    return hf_hub_download(repo_id=HF_REPO_ID, filename=filename)

scripts/FMPose3D_main.py

@@ -268,7 +268,11 @@ def print_error_action(action_error_sum, is_train):
             args.checkpoint = "./checkpoint/" + folder_name
         elif args.train == False:
             # create a new folder for the test results
-            args.previous_dir = os.path.dirname(args.model_weights_path)
+            if args.model_weights_path:
+                args.previous_dir = os.path.dirname(args.model_weights_path)
+            else:
+                # HuggingFace-downloaded weights: no local dir, use ./checkpoint/
+                args.previous_dir = "./checkpoint"
             args.checkpoint = os.path.join(args.previous_dir, folder_name)
 
         if not os.path.exists(args.checkpoint):
@@ -337,8 +341,10 @@ def print_error_action(action_error_sum, is_train):
 
     if args.reload:
         model_dict = model["CFM"].state_dict()
-        model_path = args.model_weights_path
-        print(model_path)
+        from fmpose3d.utils.weights import resolve_weights_path
+        model_path = resolve_weights_path(args.model_weights_path, args.model_type)
+
+        print(f"Loading weights from: {model_path}")
         pre_dict = torch.load(model_path, map_location=device, weights_only=True)
         for name, key in model_dict.items():
             model_dict[name] = pre_dict[name]

scripts/FMPose3D_test.sh

@@ -11,7 +11,11 @@ exp_temp=0.005
 folder_name=test_s${eval_multi_steps}_${mode}_h${num_hypothesis_list}_$(date +%Y%m%d_%H%M%S)
 
 model_type='fmpose3d_humans'
-model_weights_path='./pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
+
+# By default, weights are automatically downloaded from Hugging Face Hub.
+# To use local weights instead, uncomment the line below:
+# model_weights_path='./pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
+model_weights_path=''
 
 #Test
 python3 scripts/FMPose3D_main.py \

tests/fmpose3d_api/test_fmpose3d.py

@@ -747,8 +747,8 @@ def test_predict_raises_clear_error_without_deeplabcut(self):
         frames = np.random.randint(0, 255, (1, 64, 64, 3), dtype=np.uint8)
 
         with patch(
-            "fmpose3d.inference_api.fmpose3d.importlib.util.find_spec",
-            return_value=None,
+            "fmpose3d.inference_api.fmpose3d._require_superanimal_analyze_images",
+            side_effect=ImportError('pip install "fmpose3d[animals]"'),
         ):
             with pytest.raises(ImportError, match=r"fmpose3d\[animals\]"):
                 estimator.predict(frames)

tests/test_demo_human.py

@@ -38,7 +38,7 @@ def test_2d_pose_estimation(test_image_path, test_output_dir):
     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')
+    keypoints, scores = hrnet_pose(test_image_path, det_dim=416, num_person=1, gen_output=True, type='image')
 
     # Check output shapes
     assert keypoints is not None, "Keypoints should not be None"