diff --git a/.gitignore b/.gitignore index 784ecb0..50a8fb9 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,7 @@ Thumbs.db *.pth *.onnx *.safetensors +!empty_text_encoder/empty_embedding.safetensors # Output files output/ diff --git a/README.md b/README.md index c6a3b66..2627a8b 100644 --- a/README.md +++ b/README.md @@ -111,51 +111,33 @@ Runs pose estimation on one or more images. ### `IBB Pose — JSON to Image` -Renders an OpenPose JSON string to a pose visualisation image. +## 模型目录 ---- +- SDPose 模型：`ComfyUI/models/IBB_POSE/` +- YOLO 模型：`ComfyUI/models/yolo/` +- GroundingDINO 模型：`ComfyUI/models/grounding-dino/` -## JSON Format +首次运行时会自动下载以下资源： -Output follows the standard OpenPose JSON format: +- `teemosliang/SDPose-Body` +- `teemosliang/SDPose-Wholebody` +- GroundingDINO 配置和权重 -```json -{ - "people": [ - { - "pose_keypoints_2d": [x, y, score, ...], - "face_keypoints_2d": [...], - "hand_left_keypoints_2d": [...], - "hand_right_keypoints_2d": [...], - "foot_keypoints_2d": [...] - } - ], - "canvas_width": 512, - "canvas_height": 768 -} -``` +仓库已内置 `empty_text_encoder/empty_embedding.safetensors`，也保留了 `generate_empty_embedding.py` 以便重新生成。 -Compatible with [ComfyUI-OpenPose-Editor](https://github.com/judian17/ComfyUI-OpenPose-Editor-jd). +## 工作流 ---- +示例工作流位于： -## Skeleton Keypoint Layout +- `/workflow/ibb_pose.json` -### Body / OpenPose (18 pts) -OpenPose-18: nose · neck · shoulders · elbows · wrists · hips · knees · ankles · eyes · ears +## 兼容说明 -### WholeBody (133 pts, DWPose) -- 0–16: Body (COCO-17) -- 17–22: Feet (6) -- 23–90: Face (68) -- 91–111: Left hand (21) -- 112–132: Right hand (21) +- 依赖 ComfyUI 的 `folder_paths` / `model_management` +- Florence2 检测为可选能力，需要额外安装对应 ComfyUI 节点 +- OpenPose 编辑 JSON 可与 [ComfyUI-OpenPose-Editor-jd](https://github.com/judian17/ComfyUI-OpenPose-Editor-jd) 配合使用 ---- +## 致谢 -## Recommended Workflow - -``` -Load Image → IBB Pose Load Model → IBB Pose Run Estimation → Preview Image - └→ Save Text (JSON) -``` +- 上游实现：`judian17/ComfyUI-SDPose-OOD` +- 原始项目：`T-S-Liang/SDPose-OOD` diff --git a/__init__.py b/__init__.py index ce56305..71a7574 100644 --- a/__init__.py +++ b/__init__.py @@ -1,7 +1,3 @@ -""" -IBB_POSE - ComfyUI Pose Estimation Node -Compatible with Python 3.13, torch 2.10, CUDA 13.x -""" from .ibb_pose_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/empty_text_encoder/empty_embedding.safetensors b/empty_text_encoder/empty_embedding.safetensors new file mode 100644 index 0000000..2e213eb Binary files /dev/null and b/empty_text_encoder/empty_embedding.safetensors differ diff --git a/eval/SDPose_Evaluator.py b/eval/SDPose_Evaluator.py new file mode 100644 index 0000000..9e2d959 --- /dev/null +++ b/eval/SDPose_Evaluator.py @@ -0,0 +1,140 @@ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk +# Date: Feb. 2025 +# Description: Evaluation protocol for fine-tuned Stable Diffusion models on pose estimation benchmarks. +# License: MIT License (see LICENSE for details) + +import os +import torch +from tqdm import tqdm +from safetensors.torch import load_file + +from accelerate import Accelerator +from accelerate.logging import get_logger +from packaging import version +from tqdm.auto import tqdm + +from transformers import CLIPTextModel, CLIPTokenizer +from diffusers.utils.import_utils import is_xformers_available + +from diffusers import ( + UNet2DConditionModel, + AutoencoderKL, + DDPMScheduler +) + +from pipelines.SDPose_D_Pipeline import SDPose_D_Pipeline + +from models.HeatmapHead import get_heatmap_head +from models.ModifiedUNet import Modified_forward + +from utils.EvalMetric import get_metric +from utils.dataset_loader import prepare_dataset + + +class SDPose_Evaluator: + + def __init__( + self, + args + ): + + self.logger = get_logger(__name__, log_level="INFO") + # -------------------- Device -------------------- + + self.args = args + self.accelerator = Accelerator() + + unet = UNet2DConditionModel.from_pretrained( + self.args.checkpoint_path, subfolder="unet", revision=None, + class_embed_type="projection", projection_class_embeddings_input_dim=4, + low_cpu_mem_usage=False, device_map=None, + ) + + unet = Modified_forward(unet, keypoint_scheme = self.args.keypoint_scheme) + + vae = AutoencoderKL.from_pretrained(self.args.checkpoint_path, + subfolder='vae') + + tokenizer = CLIPTokenizer.from_pretrained(self.args.checkpoint_path, subfolder='tokenizer') + text_encoder = CLIPTextModel.from_pretrained(self.args.checkpoint_path, subfolder='text_encoder') + + # Loading Heatmap Autoencoder + dec_path = os.path.join(self.args.checkpoint_path, "decoder", "decoder.safetensors") + hm_decoder = get_heatmap_head(mode = self.args.keypoint_scheme) + hm_decoder.load_state_dict(load_file(dec_path, device="cpu"), strict=True) + + noise_scheduler = DDPMScheduler.from_pretrained( + self.args.checkpoint_path, subfolder='scheduler' + ) + + self.model = SDPose_D_Pipeline(unet = unet, vae = vae, tokenizer = tokenizer, text_encoder = text_encoder, decoder = hm_decoder, scheduler = noise_scheduler) + + # using xformers for efficient attentions. + if self.args.enable_xformers_memory_efficient_attention: + if is_xformers_available(): + import xformers + xformers_version = version.parse(xformers.__version__) + if xformers_version == version.parse("0.0.16"): + self.logger.warn( + "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details." + ) + self.model.unet.enable_xformers_memory_efficient_attention() + else: + raise ValueError("xformers is not available. Make sure it is installed correctly") + + self.model.vae.requires_grad_(False) + self.model.text_encoder.requires_grad_(False) + self.model.decoder.requires_grad_(False) + self.model.unet.requires_grad_(False) + + self.model.vae.eval() + self.model.text_encoder.eval() + self.model.decoder.eval() + self.model.unet.eval() + + @torch.no_grad() + def evaluate(self): + + val_loader = prepare_dataset(batch_size = self.args.eval_batch_size, dataset_name = self.args.dataset_name, dataset_root = self.args.dataset_root, mode = "val", num_workers = self.args.dataloader_num_workers) + + + self.model.unet, self.model.vae, self.model.decoder, val_loader, self.model.text_encoder = self.accelerator.prepare(self.model.unet, self.model.vae, self.model.decoder, val_loader, self.model.text_encoder) + + device = self.accelerator.device + + metric = get_metric(ann_file = self.args.ann_file, mode = self.args.keypoint_scheme) + + metric.dataset_meta = val_loader.dataset.metainfo + + iterator = tqdm(val_loader, desc="Validating", leave=False) if self.accelerator.is_main_process else val_loader + + for batch in iterator: + + rgb_in = batch["inputs"] # [B, 3, H, W] + rgb_in = rgb_in[:, [2, 1, 0], ...] + rgb_in = rgb_in.float() / 255.0 + + rgb_in = rgb_in * 2.0 - 1.0 + + rgb_in = rgb_in.to(device) # [B, 3, H, W] + + ds_list = batch['data_samples'] + + # Enable flip test by setting test_cfg + test_cfg = {'flip_test': True} # Re-enabled to test our fixes + + # Run inference with flip test enabled + timesteps = [self.args.timestep] + outputs = self.model(rgb_in, mode = "predict", data_samples=batch['data_samples'], timesteps = timesteps, test_cfg=test_cfg) + + samples_for_metric = [s.to_dict() for s in outputs] + + metric.process(data_samples = samples_for_metric, data_batch = batch) + + self.accelerator.wait_for_everyone() + res = metric.evaluate(len(val_loader.dataset)) + + if self.accelerator.is_main_process: + + print(res) \ No newline at end of file diff --git a/eval/launch_evaluator.py b/eval/launch_evaluator.py new file mode 100644 index 0000000..656b622 --- /dev/null +++ b/eval/launch_evaluator.py @@ -0,0 +1,82 @@ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk +# Date: Feb. 2025 +# License: MIT License (see LICENSE for details) + +from SDPose_Evaluator import SDPose_Evaluator +import argparse + +def parse_args(): + + parser = argparse.ArgumentParser(description="Evaluate SDPose on Pose Estimation Benchmarks") + + parser.add_argument( + "--dataset_name", + type=str, + default=None, + required=True, + help="The dataset name for the evaluation benchmark.", + ) + + parser.add_argument( + "--keypoint_scheme", + type=str, + default="body", + required=True, + help="The keypoint scheme for the pose estimation model, please fill in body (COCO-17) or wholebody (COCOWholebody-133).", + ) + + parser.add_argument( + "--dataset_root", + type=str, + default=None, + required=True, + help="The root path for the datasets.", + ) + + parser.add_argument( + "--checkpoint_path", + type=str, + default=None, + required=True, + help="Path to the SDPose checkpoint.", + ) + + parser.add_argument( + "--eval_batch_size", type=int, default=1, help="Batch size (per device) for the evaluation dataloader." + ) + + parser.add_argument( + "--ann_file", type=str, help="The path for the annotation file." + ) + + parser.add_argument( + "--timestep", + type=int, + default=999, + ) + + # dataloaderes + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + + # using xformers for efficient training and inference + parser.add_argument( + "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers." + ) + + args = parser.parse_args() + + return args + + +args = parse_args() +Trainer = SDPose_Evaluator(args) + +Trainer.evaluate() \ No newline at end of file diff --git a/generate_empty_embedding.py b/generate_empty_embedding.py new file mode 100644 index 0000000..8b99a3e --- /dev/null +++ b/generate_empty_embedding.py @@ -0,0 +1,67 @@ +import torch +from transformers import CLIPTokenizer, CLIPTextModel +from safetensors.torch import save_file +import os +from pathlib import Path +# No longer need folder_paths for this script + +# --- Configuration --- +# We only need to do this once, so we can hardcode the repo ID. +# The empty embedding is identical for both Body and WholeBody models. +REPO_ID = "teemosliang/SDPose-Body" +# Get the directory of the current script and build the path relative to it +SCRIPT_DIR = Path(__file__).parent +OUTPUT_DIR = os.path.join(SCRIPT_DIR, "empty_text_encoder") +OUTPUT_FILE = os.path.join(OUTPUT_DIR, "empty_embedding.safetensors") + +def generate_and_save_empty_embedding(): + """ + Loads a CLIP text model, generates an embedding for an empty prompt, + and saves it to a safetensors file. + """ + print("--- SDPose Empty Embedding Generator ---") + + # Ensure the output directory exists + os.makedirs(OUTPUT_DIR, exist_ok=True) + print(f"Output directory prepared: {OUTPUT_DIR}") + + # Check if the file already exists + if os.path.exists(OUTPUT_FILE): + print(f"'{OUTPUT_FILE}' already exists. Skipping generation.") + return + + try: + # 1. Load tokenizer and text encoder from Hugging Face + print(f"Loading CLIP model from '{REPO_ID}'...") + tokenizer = CLIPTokenizer.from_pretrained(REPO_ID, subfolder='tokenizer') + text_encoder = CLIPTextModel.from_pretrained(REPO_ID, subfolder='text_encoder') + print("Model loaded successfully.") + + # 2. Generate the empty text embedding + print("Generating empty text embedding...") + prompt = "" + text_inputs = tokenizer( + prompt, + padding="do_not_pad", + max_length=tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + empty_text_embed = text_encoder(text_input_ids)[0] + print(f"Embedding generated with shape: {empty_text_embed.shape}") + + # 3. Save the embedding to a safetensors file + print(f"Saving embedding to '{OUTPUT_FILE}'...") + tensor_dict = {"empty_text_embed": empty_text_embed} + save_file(tensor_dict, OUTPUT_FILE) + print("--- Success! ---") + print("Empty embedding has been saved. You no longer need to run this script.") + + except Exception as e: + print(f"\n--- An error occurred ---") + print(f"Error: {e}") + print("Please ensure you have an internet connection and the 'transformers' library is installed.") + +if __name__ == "__main__": + generate_and_save_empty_embedding() diff --git a/ibb_pose_nodes.py b/ibb_pose_nodes.py index 43b5183..c676184 100644 --- a/ibb_pose_nodes.py +++ b/ibb_pose_nodes.py @@ -1,25 +1,14 @@ -""" -IBB_POSE - ComfyUI Pose Estimation Node -Supports: Body (17 kpts) / WholeBody (133 kpts) / OpenPose (18 kpts) -Single-person & multi-person detection, pose image + JSON output -Compatible with Python 3.13, torch 2.10+, CUDA 12/13 -""" - import warnings import logging -import os -import gc -import sys -import json -import math -import colorsys -import urllib.request -import traceback -import time -from pathlib import Path +# Suppress common warnings to reduce noise warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) +warnings.filterwarnings("ignore", message=".*Redirects are currently not supported.*") + +# Set logging level to reduce verbose output +logging.getLogger("torch.distributed").setLevel(logging.ERROR) +logging.getLogger("deepspeed").setLevel(logging.ERROR) def _exc_summary(exc: Exception) -> str: return f"{exc.__class__.__name__}: {exc}" @@ -65,54 +54,84 @@ except Exception as _e: ) import numpy as np -import cv2 from PIL import Image - -# ComfyUI integration (graceful fallback for standalone testing) +import json try: import folder_paths - IS_COMFYUI = True except ImportError: - class _FolderPaths: - models_dir = os.path.join(os.path.expanduser("~"), "ComfyUI", "models") - output_dir = os.path.join(os.path.expanduser("~"), "ComfyUI", "output") + # If not available, create a mock object for testing + class MockFolderPaths: + models_dir = os.path.expanduser("~/models") + output_dir = os.path.expanduser("~/output") supported_pt_extensions = {".pt", ".pth", ".ckpt", ".safetensors"} - folder_names_and_paths: dict = {} + folder_names_and_paths = {} - @staticmethod - def get_filename_list(folder_name): - return [] + @classmethod + def add_model_folder_path(cls, folder_name, full_folder_path, is_default=True): + cls.folder_names_and_paths.setdefault(folder_name, []) + if full_folder_path not in cls.folder_names_and_paths[folder_name]: + cls.folder_names_and_paths[folder_name].append(full_folder_path) - @staticmethod - def get_full_path(folder_name, filename): + @classmethod + def get_folder_paths(cls, folder_name): + return cls.folder_names_and_paths.get(folder_name, []) + + @classmethod + def get_filename_list(cls, folder_name): + names = [] + for base in cls.get_folder_paths(folder_name): + if not os.path.isdir(base): + continue + for entry in sorted(os.listdir(base)): + path = os.path.join(base, entry) + if os.path.isfile(path): + names.append(entry) + return names + + @classmethod + def get_full_path(cls, folder_name, filename): + for base in cls.get_folder_paths(folder_name): + path = os.path.join(base, filename) + if os.path.exists(path): + return path return None - @staticmethod - def add_model_folder_path(folder_name, full_folder_path, is_default=True): - pass - - @staticmethod - def get_output_directory(): - return _FolderPaths.output_dir - - folder_paths = _FolderPaths() - IS_COMFYUI = False + @classmethod + def get_output_directory(cls): + os.makedirs(cls.output_dir, exist_ok=True) + return cls.output_dir + @classmethod + def get_save_image_path(cls, filename_prefix, output_dir, width, height): + subfolder = os.path.dirname(filename_prefix) + base_filename = os.path.basename(filename_prefix) + full_output_folder = os.path.join(output_dir, subfolder) if subfolder else output_dir + os.makedirs(full_output_folder, exist_ok=True) + return full_output_folder, base_filename, 0, subfolder, filename_prefix + folder_paths = MockFolderPaths() +from huggingface_hub import snapshot_download +import sys +from pathlib import Path +import glob +import cv2 +import math +import matplotlib.colors +import tempfile +from torchvision import transforms try: import model_management except ImportError: - class _ModelManagement: + # If not available, create a mock object for testing + class MockModelManagement: @staticmethod def get_torch_device(): return torch.device("cuda" if torch.cuda.is_available() else "cpu") - @staticmethod def soft_empty_cache(): if torch.cuda.is_available(): torch.cuda.empty_cache() - gc.collect() + model_management = MockModelManagement() - model_management = _ModelManagement() # Optional dependencies try: @@ -137,12 +156,13 @@ except Exception as _e: "WholeBody mode disabled." ) -# Model directory -IBB_POSE_MODEL_DIR = os.path.join(folder_paths.models_dir, "IBB_POSE") -os.makedirs(IBB_POSE_MODEL_DIR, exist_ok=True) +# --- Add custom folder paths to ComfyUI --- +SDPOSE_MODEL_DIR = os.path.join(folder_paths.models_dir, "IBB_POSE") +YOLO_MODEL_DIR = os.path.join(folder_paths.models_dir, "yolo") try: - folder_paths.add_model_folder_path("IBB_POSE", IBB_POSE_MODEL_DIR) + folder_paths.add_model_folder_path("IBB_POSE", SDPOSE_MODEL_DIR, is_default=True) + folder_paths.add_model_folder_path("yolo", YOLO_MODEL_DIR, is_default=False) except Exception: pass @@ -355,320 +375,475 @@ def _draw_body(canvas, op_kpts, op_scores, threshold=0.3, pose_scale=1.0): return canvas +def draw_wholebody_keypoints_openpose_style(canvas, keypoints, scores=None, threshold=0.3, overlay_mode=False, overlay_alpha=0.6, scale_for_xinsr=False, pose_scale=1.0): + H, W, C = canvas.shape + + # --- 动态计算粗细 (基于 pose_scale) --- + base_stickwidth = int(4 * pose_scale) + base_radius = int(4 * pose_scale) + + face_hand_radius = max(1, int(3 * pose_scale)) + face_hand_line_width = max(1, int(2 * pose_scale)) -def _draw_wholebody(canvas, kpts, scores, threshold=0.3, - keep_face=True, keep_hands=True, keep_feet=True, pose_scale=1.0): - H, W = canvas.shape[:2] - avg = (H + W) / 2.0 - sw = max(1, int(4 * pose_scale)) - cr = max(1, int(4 * pose_scale)) - f_r = max(1, int(3 * pose_scale)) - f_lw = max(1, int(2 * pose_scale)) + stickwidth = base_stickwidth - for i, (a1, b1) in enumerate(OPENPOSE_LIMB_SEQ): - a, b = a1 - 1, b1 - 1 - if a >= len(scores) or b >= len(scores): - continue - if scores[a] < threshold or scores[b] < threshold: - continue - ax, ay = float(kpts[a, 0]), float(kpts[a, 1]) - bx, by = float(kpts[b, 0]), float(kpts[b, 1]) - mX = (ax + bx) / 2.0 - mY = (ay + by) / 2.0 - length = math.hypot(ax - bx, ay - by) - if length < 1.0: - continue - angle = math.degrees(math.atan2(ax - bx, ay - by)) - polygon = cv2.ellipse2Poly( - (int(mX), int(mY)), (max(1, int(length / 2)), sw), - int(angle), 0, 360, 1 - ) - cv2.fillConvexPoly(canvas, polygon, BODY_COLORS[i % len(BODY_COLORS)]) + # --- Xinsr Logic (Fixed Multiplier) --- + if scale_for_xinsr: + xinsr_fixed_multiplier = 2 + stickwidth = int(base_stickwidth * xinsr_fixed_multiplier) - for i in range(min(18, len(kpts))): - if i >= len(scores) or scores[i] < threshold: - continue - x, y = int(kpts[i, 0]), int(kpts[i, 1]) - if 0 <= x < W and 0 <= y < H: - cv2.circle(canvas, (x, y), cr, BODY_COLORS[i % len(BODY_COLORS)], -1) - - if keep_feet and len(kpts) >= 24: + body_limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], [1, 16], [16, 18]] + hand_edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] + colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] + + # Body Limbs + if len(keypoints) >= 18: + for i, limb in enumerate(body_limbSeq): + idx1, idx2 = limb[0] - 1, limb[1] - 1 + if idx1 >= 18 or idx2 >= 18: continue + if scores is not None and (scores[idx1] < threshold or scores[idx2] < threshold): continue + Y = np.array([keypoints[idx1][0], keypoints[idx2][0]]); X = np.array([keypoints[idx1][1], keypoints[idx2][1]]); mX = np.mean(X); mY = np.mean(Y) + length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 + if length < 1: continue + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, colors[i % len(colors)]) + for i in range(18): + if scores is not None and scores[i] < threshold: continue + x, y = int(keypoints[i][0]), int(keypoints[i][1]) + if 0 <= x < W and 0 <= y < H: cv2.circle(canvas, (x, y), base_radius, colors[i % len(colors)], thickness=-1) + + # Feet + if len(keypoints) >= 24: for i in range(18, 24): - if i >= len(scores) or scores[i] < threshold: - continue - x, y = int(kpts[i, 0]), int(kpts[i, 1]) - if 0 <= x < W and 0 <= y < H: - cv2.circle(canvas, (x, y), cr, BODY_COLORS[i % len(BODY_COLORS)], -1) - - if keep_face and len(kpts) >= 92: + if scores is not None and scores[i] < threshold: continue + x, y = int(keypoints[i][0]), int(keypoints[i][1]) + if 0 <= x < W and 0 <= y < H: cv2.circle(canvas, (x, y), base_radius, colors[i % len(colors)], thickness=-1) + + # Hands + if len(keypoints) >= 113: + for ie, edge in enumerate(hand_edges): + idx1, idx2 = 92 + edge[0], 92 + edge[1] + if scores is not None and (scores[idx1] < threshold or scores[idx2] < threshold): continue + x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1]); x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1]) + if x1 > 0.01 and y1 > 0.01 and x2 > 0.01 and y2 > 0.01 and 0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H: + color = matplotlib.colors.hsv_to_rgb([ie / float(len(hand_edges)), 1.0, 1.0]) * 255 + cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=face_hand_line_width) + for i in range(92, 113): + if scores is not None and scores[i] < threshold: continue + x, y = int(keypoints[i][0]), int(keypoints[i][1]) + if x > 0.01 and y > 0.01 and 0 <= x < W and 0 <= y < H: cv2.circle(canvas, (x, y), face_hand_radius, (0, 0, 255), thickness=-1) + if len(keypoints) >= 134: + for ie, edge in enumerate(hand_edges): + idx1, idx2 = 113 + edge[0], 113 + edge[1] + if scores is not None and (scores[idx1] < threshold or scores[idx2] < threshold): continue + x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1]); x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1]) + if x1 > 0.01 and y1 > 0.01 and x2 > 0.01 and y2 > 0.01 and 0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H: + color = matplotlib.colors.hsv_to_rgb([ie / float(len(hand_edges)), 1.0, 1.0]) * 255 + cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=face_hand_line_width) + for i in range(113, 134): + if scores is not None and i < len(scores) and scores[i] < threshold: continue + x, y = int(keypoints[i][0]), int(keypoints[i][1]) + if x > 0.01 and y > 0.01 and 0 <= x < W and 0 <= y < H: cv2.circle(canvas, (x, y), face_hand_radius, (0, 0, 255), thickness=-1) + + # Face + if len(keypoints) >= 92: for i in range(24, 92): - if i >= len(scores) or scores[i] < threshold: - continue - x, y = int(kpts[i, 0]), int(kpts[i, 1]) - if x > 1 and y > 1 and 0 <= x < W and 0 <= y < H: - cv2.circle(canvas, (x, y), f_r, (255, 255, 255), -1) - - def _draw_hand(start_idx): - if len(kpts) < start_idx + 21: - return - for ie, (ea, eb) in enumerate(HAND_EDGES): - a, b = start_idx + ea, start_idx + eb - if a >= len(scores) or b >= len(scores): - continue - if scores[a] < threshold or scores[b] < threshold: - continue - x1, y1 = int(kpts[a, 0]), int(kpts[a, 1]) - x2, y2 = int(kpts[b, 0]), int(kpts[b, 1]) - if (x1 > 1 and y1 > 1 and x2 > 1 and y2 > 1 - and 0 <= x1 < W and 0 <= y1 < H - and 0 <= x2 < W and 0 <= y2 < H): - red, green, blue = colorsys.hsv_to_rgb(ie / len(HAND_EDGES), 1.0, 1.0) - color = (int(blue * 255), int(green * 255), int(red * 255)) - cv2.line(canvas, (x1, y1), (x2, y2), color, f_lw) - for i in range(start_idx, start_idx + 21): - if i >= len(scores) or scores[i] < threshold: - continue - x, y = int(kpts[i, 0]), int(kpts[i, 1]) - if x > 1 and y > 1 and 0 <= x < W and 0 <= y < H: - cv2.circle(canvas, (x, y), f_r, (0, 0, 255), -1) - - if keep_hands: - _draw_hand(92) - _draw_hand(113) - + if scores is not None and scores[i] < threshold: continue + x, y = int(keypoints[i][0]), int(keypoints[i][1]) + if x > 0.01 and y > 0.01 and 0 <= x < W and 0 <= y < H: cv2.circle(canvas, (x, y), face_hand_radius, (255, 255, 255), thickness=-1) return canvas - -def _fmt_kpts(kpts, scores, start, end, threshold, filter_low): - out = [] - for i in range(start, end): - if i >= len(kpts): - out.extend([0.0, 0.0, 0.0]) - else: - s = float(scores[i]) if i < len(scores) else 0.0 - if filter_low and s < threshold: - out.extend([0.0, 0.0, 0.0]) - else: - out.extend([float(kpts[i, 0]), float(kpts[i, 1]), s]) - return out - - -def _build_openpose_json(all_kpts, all_scores, image_width, image_height, - skeleton_type="body", threshold=0.3, filter_low=True): - if image_width == 0 or image_height == 0: - return json.dumps({"people": [], "canvas_width": image_width, - "canvas_height": image_height}) - - people = [] - for kpts, scores in zip(all_kpts, all_scores): - if kpts is None or len(kpts) == 0: - continue - person = {} - - if skeleton_type in ("body", "openpose"): - if kpts.shape[0] == 17: - op_k, op_s = _coco17_to_openpose18(kpts, scores) - else: - op_k, op_s = kpts[:18], scores[:18] - person["pose_keypoints_2d"] = _fmt_kpts(op_k, op_s, 0, 18, threshold, filter_low) - person["face_keypoints_2d"] = [] - person["hand_left_keypoints_2d"] = [] - person["hand_right_keypoints_2d"] = [] - person["foot_keypoints_2d"] = [] - else: - person["pose_keypoints_2d"] = _fmt_kpts(kpts, scores, 0, 18, threshold, filter_low) - person["foot_keypoints_2d"] = _fmt_kpts(kpts, scores, 18, 24, threshold, filter_low) - person["face_keypoints_2d"] = _fmt_kpts(kpts, scores, 24, 92, threshold, filter_low) - person["hand_left_keypoints_2d"] = _fmt_kpts(kpts, scores, 92, 113, threshold, filter_low) - person["hand_right_keypoints_2d"] = _fmt_kpts(kpts, scores, 113, 134, threshold, filter_low) - - people.append(person) - - return json.dumps({"people": people, - "canvas_width": int(image_width), - "canvas_height": int(image_height)}, indent=2) - - -def _preprocess_for_dwpose(img_bgr, input_wh=(288, 384)): - W_in, H_in = input_wh - img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) - img_resized = cv2.resize(img_rgb, (W_in, H_in)).astype(np.float32) / 255.0 - mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) - std = np.array([0.229, 0.224, 0.225], dtype=np.float32) - img_norm = (img_resized - mean) / std - return img_norm.transpose(2, 0, 1)[np.newaxis].astype(np.float32) - - -def _decode_simcc(simcc_x, simcc_y, split_ratio=2.0): - kpts_x = np.argmax(simcc_x[0], axis=1).astype(np.float32) / split_ratio - kpts_y = np.argmax(simcc_y[0], axis=1).astype(np.float32) / split_ratio - score_x = np.max(simcc_x[0], axis=1) - score_y = np.max(simcc_y[0], axis=1) - scores = (score_x + score_y) / 2.0 - kpts = np.stack([kpts_x, kpts_y], axis=1) - return kpts, scores - - -def _run_dwpose_crop(pose_session, img_bgr, bbox, input_wh=(288, 384)): - H_img, W_img = img_bgr.shape[:2] - - if bbox is not None: - x1 = max(0, int(bbox[0])) - y1 = max(0, int(bbox[1])) - x2 = min(W_img, int(bbox[2])) - y2 = min(H_img, int(bbox[3])) - if x2 <= x1 or y2 <= y1: - return None, None - crop = img_bgr[y1:y2, x1:x2] - crop_w = x2 - x1 - crop_h = y2 - y1 - offset_x = x1 - offset_y = y1 - else: - crop = img_bgr - crop_w = W_img - crop_h = H_img - offset_x = 0 - offset_y = 0 - - W_in, H_in = input_wh - inp = _preprocess_for_dwpose(crop, input_wh) - +# --- GroundingDINO Prediction Functions (Adapted from SAM2 Node) --- +# (modified for SDPose context) +def load_dino_image(image_pil): + # Need import from local_groundingdino here try: - iname = pose_session.get_inputs()[0].name - output = pose_session.run(None, {iname: inp}) - except Exception as exc: - print(f"IBB_POSE: DWPose inference error: {exc}") - return None, None + from groundingdino.datasets import transforms as T + except ImportError: + raise ImportError("IBB_POSE: Failed to import 'groundingdino'. Please install it via pip: 'pip install groundingdino-py'") - simcc_x, simcc_y = output[0], output[1] - kpts_crop, scores = _decode_simcc(simcc_x, simcc_y) - - kpts_crop[:, 0] = kpts_crop[:, 0] * (crop_w / W_in) + offset_x - kpts_crop[:, 1] = kpts_crop[:, 1] * (crop_h / H_in) + offset_y - - kpts_out, scores_out = _coco_wholebody_reorder(kpts_crop, scores) - return kpts_out, scores_out - - -def _letterbox_yolox(img, new_shape=(640, 640)): - h0, w0 = img.shape[:2] - r = min(new_shape[0] / h0, new_shape[1] / w0) - h1 = int(round(h0 * r)) - w1 = int(round(w0 * r)) - dh = (new_shape[0] - h1) / 2.0 - dw = (new_shape[1] - w1) / 2.0 - img_r = cv2.resize(img, (w1, h1)) - top = int(round(dh - 0.1)) - bottom = int(round(dh + 0.1)) - left = int(round(dw - 0.1)) - right = int(round(dw + 0.1)) - padded = cv2.copyMakeBorder( - img_r, top, bottom, left, right, - cv2.BORDER_CONSTANT, value=(114, 114, 114) + transform = T.Compose( + [ + T.RandomResize([800], max_size=1333), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] ) - return padded, r, (dw, dh) + image, _ = transform(image_pil, None) # 3, h, w + return image +# (modified for SDPose context) +def groundingdino_predict(dino_model_wrapper, image_pil, prompt, threshold): + dino_model = dino_model_wrapper["model"] # Extract the actual model + + # (nested function) + def get_grounding_output(model, image, caption, box_threshold, device): + caption = caption.lower().strip() + if not caption.endswith("."): + caption = caption + "." + + # Move model to device just before inference + model.to(device) + image = image.to(device) + + boxes_filt = torch.tensor([]) # Initialize empty tensor + try: + with torch.no_grad(): + outputs = model(image[None], captions=[caption]) + + logits = outputs["pred_logits"].sigmoid()[0] # (nq, 256) + boxes = outputs["pred_boxes"][0] # (nq, 4) in cxcywh format + + # Filter output based on threshold + filt_mask = logits.max(dim=1)[0] > box_threshold + logits_filt = logits[filt_mask] # num_filt, 256 + boxes_filt = boxes[filt_mask] # num_filt, 4 + + # Convert boxes from cxcywh to xyxy format + boxes_filt = box_cxcywh_to_xyxy(boxes_filt) -def _nms(boxes, scores, iou_thresh=0.45): - idxs = np.argsort(scores)[::-1] - keep = [] - while len(idxs): - i = int(idxs[0]) - keep.append(i) - if len(idxs) == 1: - break - rest = idxs[1:] - xi1 = np.maximum(boxes[i, 0], boxes[rest, 0]) - yi1 = np.maximum(boxes[i, 1], boxes[rest, 1]) - xi2 = np.minimum(boxes[i, 2], boxes[rest, 2]) - yi2 = np.minimum(boxes[i, 3], boxes[rest, 3]) - inter = np.maximum(xi2 - xi1, 0) * np.maximum(yi2 - yi1, 0) - ai = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1]) - ar = (boxes[rest, 2] - boxes[rest, 0]) * (boxes[rest, 3] - boxes[rest, 1]) - union = ai + ar - inter - iou = np.where(union > 0, inter / union, 0.0) - idxs = rest[iou < iou_thresh] - return keep + except Exception as e: + print(f"IBB_POSE: GroundingDINO inference failed: {e}") + finally: + # Move model back to CPU after inference to save VRAM + model.to("cpu") + model_management.soft_empty_cache() # Clean VRAM + + return boxes_filt.cpu() # Return results on CPU + # Helper function to convert box format + def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.stack(b, dim=1) + + # --- Main prediction logic --- + dino_image = load_dino_image(image_pil.convert("RGB")) + + # Use ComfyUI's device management + device = model_management.get_torch_device() + + boxes_filt_norm = get_grounding_output(dino_model, dino_image, prompt, threshold, device) + + # If no boxes found, return empty list + if boxes_filt_norm.shape[0] == 0: + return [] + + # Scale boxes to original image size + H, W = image_pil.size[1], image_pil.size[0] + boxes_filt_abs = boxes_filt_norm * torch.Tensor([W, H, W, H]) + + # Convert to list of lists format [[x1, y1, x2, y2], ...] + return boxes_filt_abs.tolist() -def _detect_persons_yolox(det_session, img_bgr, conf_thresh=0.45): - H_orig, W_orig = img_bgr.shape[:2] - fallback = [[0, 0, W_orig, H_orig]] - - padded, scale, (dw, dh) = _letterbox_yolox(img_bgr, (640, 640)) - padded_rgb = padded[:, :, ::-1].astype(np.float32) - inp = padded_rgb.transpose(2, 0, 1)[np.newaxis] - +# --- Processing functions (adapted from SDPose_gradio.py) --- +# (Functions detect_person_yolo, preprocess_image_for_sdpose, restore_keypoints_to_original, convert_to_openpose_json are omitted for brevity but would be pasted here) +def detect_person_yolo(image, yolo_model_path, confidence_threshold=0.5): + if not YOLO_AVAILABLE: return [[0, 0, image.shape[1], image.shape[0]]], False try: - iname = det_session.get_inputs()[0].name - preds = det_session.run(None, {iname: inp})[0] - preds = preds[0] - except Exception as exc: - print(f"IBB_POSE: YOLOX detection error: {exc}") - return fallback + model = YOLO(yolo_model_path) + results = model(image, verbose=False) + person_bboxes = [] + for result in results: + if result.boxes is not None: + for box in result.boxes: + if int(box.cls[0]) == 0 and float(box.conf[0]) > confidence_threshold: + person_bboxes.append(box.xyxy[0].cpu().numpy().tolist()) + if person_bboxes: return person_bboxes, True + else: return [[0, 0, image.shape[1], image.shape[0]]], False + except Exception as e: + print(f"IBB_POSE: YOLO detection failed: {e}") + return [[0, 0, image.shape[1], image.shape[0]]], False - if preds.ndim != 2 or preds.shape[1] < 6: - return fallback - - obj_conf = preds[:, 4] - cls_score = preds[:, 5:] - person_s = cls_score[:, 0] * obj_conf - mask = person_s > conf_thresh - - if not np.any(mask): - return fallback - - boxes_cxcywh = preds[mask, :4] - scores_f = person_s[mask] - - boxes_xyxy = np.zeros_like(boxes_cxcywh) - boxes_xyxy[:, 0] = boxes_cxcywh[:, 0] - boxes_cxcywh[:, 2] / 2 - boxes_xyxy[:, 1] = boxes_cxcywh[:, 1] - boxes_cxcywh[:, 3] / 2 - boxes_xyxy[:, 2] = boxes_cxcywh[:, 0] + boxes_cxcywh[:, 2] / 2 - boxes_xyxy[:, 3] = boxes_cxcywh[:, 1] + boxes_cxcywh[:, 3] / 2 - - keep = _nms(boxes_xyxy, scores_f) - boxes_xyxy = boxes_xyxy[keep] - - result = [] - for box in boxes_xyxy: - x1 = max(0, int((box[0] - dw) / scale)) - y1 = max(0, int((box[1] - dh) / scale)) - x2 = min(W_orig, int((box[2] - dw) / scale)) - y2 = min(H_orig, int((box[3] - dh) / scale)) +def preprocess_image_for_sdpose(image, bbox=None, input_size=(768, 1024)): + if isinstance(image, np.ndarray): + pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) + else: pil_image = image + original_size = pil_image.size + crop_info = (0, 0, pil_image.width, pil_image.height) + if bbox is not None: + x1, y1, x2, y2 = map(int, bbox) if x2 > x1 and y2 > y1: - result.append([x1, y1, x2, y2]) + pil_image = pil_image.crop((x1, y1, x2, y2)) + crop_info = (x1, y1, x2 - x1, y2 - y1) + + transform = transforms.Compose([ + transforms.Resize((input_size[1], input_size[0]), interpolation=transforms.InterpolationMode.BILINEAR), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + ]) + input_tensor = transform(pil_image).unsqueeze(0) + return input_tensor, original_size, crop_info - return result if result else fallback +def restore_keypoints_to_original(keypoints, crop_info, input_size, original_size): + x1, y1, crop_w, crop_h = crop_info + scale_x = crop_w / input_size[0]; scale_y = crop_h / input_size[1] + keypoints_restored = keypoints.copy() + keypoints_restored[:, 0] = keypoints[:, 0] * scale_x + x1 + keypoints_restored[:, 1] = keypoints[:, 1] * scale_y + y1 + return keypoints_restored +def convert_to_loader_json(all_keypoints, all_scores, image_width, image_height, keypoint_scheme="body", threshold=0.3, enable_filter=True): -class IBBLoadPoseModel: """ - IBB_POSE — Load Model - - Body / OpenPose: ultralytics YOLO pose (17 COCO keypoints -> OpenPose-18). - WholeBody: DWPose ONNX (133 keypoints: body + face + hands + feet). - auto_download: Fetch models on first use; disable for air-gapped systems. + 将关键点转换为编辑器可读的 "Loader" JSON 格式 (OpenPose-18 body only)。 + 无论输入是 Body (17点) 还是 WholeBody (134点)，都统一输出为 OpenPose 18点格式。 """ + people = [] + + # 确保输入是列表 + if not isinstance(all_keypoints, list): + return {"width": int(image_width), "height": int(image_height), "people": []} + for keypoints, scores in zip(all_keypoints, all_scores): + person_data = {} + pose_kpts_18 = [] + + # 确保 keypoints 是 numpy 数组 + if not isinstance(keypoints, np.ndarray): + keypoints = np.array(keypoints) + if not isinstance(scores, np.ndarray): + scores = np.array(scores) + + if keypoint_scheme == "body": + # --- Body (COCO-17) -> 需要计算 Neck 转换为 OpenPose-18 --- + if len(keypoints) < 17 or len(scores) < 17: + continue + + # 1. 计算 Neck (Shoulder 中点) + neck = (keypoints[5] + keypoints[6]) / 2 + neck_score = min(scores[5], scores[6]) + + # 2. 映射关系 + op_keypoints = np.zeros((18, 2)) + op_scores = np.zeros(18) + + coco_to_op_map = [0, -1, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3] + + for i_op in range(18): + if i_op == 1: # Neck + op_keypoints[i_op] = neck + op_scores[i_op] = neck_score + else: + i_coco = coco_to_op_map[i_op] + if i_coco < len(keypoints): + op_keypoints[i_op] = keypoints[i_coco] + op_scores[i_op] = scores[i_coco] + + # 3. 扁平化 [x, y, score] + for i in range(18): + score = float(op_scores[i]) + if enable_filter and score < threshold: + pose_kpts_18.extend([0.0, 0.0, 0.0]) + else: + pose_kpts_18.extend([float(op_keypoints[i, 0]), float(op_keypoints[i, 1]), score]) + + else: + # --- WholeBody (Already OpenPose format) --- + # 这里的 keypoints 已经在 process_sequence 中被重组为 OpenPose 格式 + # 所以前 18 个点就是我们需要的 Body 18 + if len(keypoints) < 18 or len(scores) < 18: + continue + + for i in range(18): + score = float(scores[i]) + if enable_filter and score < threshold: + pose_kpts_18.extend([0.0, 0.0, 0.0]) + else: + pose_kpts_18.extend([float(keypoints[i, 0]), float(keypoints[i, 1]), score]) + + person_data["pose_keypoints_2d"] = pose_kpts_18 + people.append(person_data) + + return {"width": int(image_width), "height": int(image_height), "people": people} + +def convert_to_openpose_json(all_keypoints, all_scores, image_width, image_height, keypoint_scheme="body"): + people = [] + + # 防止除以零错误 + if image_width == 0 or image_height == 0: + return {"people": [], "canvas_width": int(image_width), "canvas_height": int(image_height)} + + for person_idx, (keypoints, scores) in enumerate(zip(all_keypoints, all_scores)): + person_data = {} + + # --- 修复核心：恢复归一化 (Normalized Coordinates) --- + # OpenPose 预览工具 (util.py) 期望输入是 0.0-1.0，它会在绘图时自己乘以 W/H。 + def format_kpt(x, y, s): + return [ + float(x), + float(y), + float(s) + ] + + if keypoint_scheme == "body": + # --- Body Only (17 -> 18) --- + if len(keypoints) < 17: continue + + # 计算颈部 (Neck) + neck = (keypoints[5] + keypoints[6]) / 2 + neck_score = min(scores[5], scores[6]) + + # 映射关系 COCO -> OpenPose Body 18 + coco_to_op_map = [0, -1, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3] + + pose_kpts_18 = [] + for i_op in range(18): + if i_op == 1: # Neck + pose_kpts_18.extend(format_kpt(neck[0], neck[1], neck_score)) + else: + i_coco = coco_to_op_map[i_op] + if i_coco < len(keypoints): + pose_kpts_18.extend(format_kpt(keypoints[i_coco, 0], keypoints[i_coco, 1], scores[i_coco])) + else: + pose_kpts_18.extend([0.0, 0.0, 0.0]) + + person_data["pose_keypoints_2d"] = pose_kpts_18 + person_data["face_keypoints_2d"] = [] + person_data["hand_left_keypoints_2d"] = [] + person_data["hand_right_keypoints_2d"] = [] + person_data["foot_keypoints_2d"] = [] + + else: + # --- WholeBody (134 Points) --- + # 这里的 keypoints 是 process_sequence 处理后的： + # 0-17: Body (含插入的 Neck) + # 18-23: Foot + # 24-91: Face (68点) + # 92-112: Left Hand (21点) + # 113-133: Right Hand (21点) + + # 1. Body (0-18) + pose_kpts = [] + for i in range(18): + pose_kpts.extend(format_kpt(keypoints[i, 0], keypoints[i, 1], scores[i])) + + # 2. Foot (18-24) + foot_kpts = [] + if len(keypoints) >= 24: + for i in range(18, 24): + foot_kpts.extend(format_kpt(keypoints[i, 0], keypoints[i, 1], scores[i])) + + # 3. Face (24-92) -> 70 点 (68 面部标志 + 双眼) + face_kpts = [] + if len(keypoints) >= 92: + # 68 面部标志点 + for i in range(24, 92): + face_kpts.extend(format_kpt(keypoints[i, 0], keypoints[i, 1], scores[i])) + # 添加右眼 (body[14]) 和左眼 (body[15]) + face_kpts.extend(format_kpt(keypoints[14, 0], keypoints[14, 1], scores[14])) + face_kpts.extend(format_kpt(keypoints[15, 0], keypoints[15, 1], scores[15])) + + # 4. Right Hand (92-113) - 官方规范：92-112 = 右手 + right_hand_kpts = [] + if len(keypoints) >= 113: + for i in range(92, 113): + right_hand_kpts.extend(format_kpt(keypoints[i, 0], keypoints[i, 1], scores[i])) + + # 5. Left Hand (113-134) - 官方规范：113-133 = 左手 + left_hand_kpts = [] + if len(keypoints) >= 134: + for i in range(113, 134): + left_hand_kpts.extend(format_kpt(keypoints[i, 0], keypoints[i, 1], scores[i])) + + person_data["pose_keypoints_2d"] = pose_kpts + person_data["face_keypoints_2d"] = face_kpts + person_data["hand_right_keypoints_2d"] = right_hand_kpts + person_data["hand_left_keypoints_2d"] = left_hand_kpts + person_data["foot_keypoints_2d"] = foot_kpts + + people.append(person_data) + + return {"people": people, "canvas_width": int(image_width), "canvas_height": int(image_height)} + +def _combine_frame_jsons(frame_jsons): + """合并所有帧的JSON数据""" + combined_data = { + "frame_count": len(frame_jsons), + "frames": [] + } + + for i, frame_json_str in enumerate(frame_jsons): + try: + # 检查 frame_json_str 是否已经是字典 + if isinstance(frame_json_str, dict): + frame_data = frame_json_str + else: + frame_data = json.loads(frame_json_str) + + frame_data["frame_index"] = i + combined_data["frames"].append(frame_data) + except json.JSONDecodeError as e: + print(f"IBB_POSE: Warning - Failed to parse JSON for frame {i}: {e}") + combined_data["frames"].append({"frame_index": i, "error": "JSON parsing failed"}) + + return json.dumps(combined_data, indent=2) + +class IBBYOLOModelLoader: + """ComfyUI node to load a YOLO model for object detection.""" @classmethod - def INPUT_TYPES(cls): + def INPUT_TYPES(s): return { "required": { - "skeleton_type": (["Body", "WholeBody", "OpenPose"],), - "model_size": (["small", "medium", "large"],), - "device": (["auto", "cuda", "cpu"],), - "precision": (["fp32", "fp16", "bf16"],), - "auto_download": ("BOOLEAN", {"default": True}), + "model_name": (folder_paths.get_filename_list("yolo"),), } } - RETURN_TYPES = ("IBB_POSE_MODEL",) - RETURN_NAMES = ("ibb_pose_model",) - FUNCTION = "load_model" - CATEGORY = "IBB_POSE" - OUTPUT_NODE = False + RETURN_TYPES = ("YOLO_MODEL",) + FUNCTION = "load_yolo_model" + CATEGORY = "IBB_POSE" + + def load_yolo_model(self, model_name): + if not YOLO_AVAILABLE: + raise ImportError("ultralytics library is not available. Please install it to use YOLO models.") + + model_path = folder_paths.get_full_path("yolo", model_name) + if not model_path: + raise FileNotFoundError(f"YOLO model not found: {model_name}") + + print(f"IBB_POSE: Loading YOLO model from {model_path}") + model = YOLO(model_path) + return (model,) + + +class IBBPoseModelLoader: + """ComfyUI node to load SDPose models and dependencies.""" + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_type": (["Body", "WholeBody"],), + "unet_precision": (["fp32", "fp16", "bf16"],), + "device": (["auto", "cuda", "cpu"],), + "unload_on_finish": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IBB_POSE_MODEL",) + FUNCTION = "load_pose_model" + CATEGORY = "IBB_POSE" + + def get_model_path(self, repo_name): + model_pathes = folder_paths.get_folder_paths("IBB_POSE") + for path in model_pathes: + model_path = os.path.join(path, repo_name) + if os.path.exists(os.path.join(model_path, "unet")): + return model_path + + return os.path.join(SDPOSE_MODEL_DIR, repo_name) + + def load_pose_model(self, model_type, unet_precision, device, unload_on_finish): + repo_id = { + "Body": "teemosliang/SDPose-Body", + "WholeBody": "teemosliang/SDPose-Wholebody" + }[model_type] + + keypoint_scheme = model_type.lower() + model_path = self.get_model_path(repo_id.split('/')[-1]) + if not os.path.exists(os.path.join(model_path, "unet")): + print(f"IBB_POSE: Downloading model from {repo_id} to {model_path}") + snapshot_download(repo_id=repo_id, local_dir=model_path, local_dir_use_symlinks=False) def load_model(self, skeleton_type, model_size, device, precision, auto_download): if not TORCH_AVAILABLE: @@ -678,22 +853,63 @@ class IBBLoadPoseModel: ) if device == "auto": - dev = model_management.get_torch_device() + device = model_management.get_torch_device() else: - dev = torch.device(device) + device = torch.device(device) - dtype_map = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16} - dtype = dtype_map[precision] - if dev.type == "cpu" and dtype != torch.float32: - print("IBB_POSE: CPU only supports fp32 – switching to fp32.") + # --- Precision Handling --- + dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[unet_precision] + if device.type == 'cpu' and dtype == torch.float16: + print("IBB_POSE Warning: FP16 is not supported on CPU. Falling back to FP32.") dtype = torch.float32 - model_info = { - "skeleton_type": skeleton_type.lower(), - "device": dev, - "dtype": dtype, - "auto_download": auto_download, + print(f"IBB_POSE: Loading model on device: {device} with UNet precision: {unet_precision}") + + # --- Load Empty Embedding --- + embed_path = os.path.join(EMPTY_EMBED_DIR, "empty_embedding.safetensors") + if not os.path.exists(embed_path): + raise FileNotFoundError(f"Empty embedding not found at '{embed_path}'. Please run 'generate_empty_embedding.py' script first.") + empty_text_embed = load_file(embed_path)["empty_text_embed"].to(device) + + # --- Load Models --- + unet = UNet2DConditionModel.from_pretrained(model_path, subfolder="unet", torch_dtype=dtype).to(device) + unet = Modified_forward(unet, keypoint_scheme=keypoint_scheme) + vae = AutoencoderKL.from_pretrained(model_path, subfolder='vae').to(device) # VAE is small, keep fp32 + + dec_path = os.path.join(model_path, "decoder", "decoder.safetensors") + hm_decoder = get_heatmap_head(mode=keypoint_scheme).to(device) + if os.path.exists(dec_path): + hm_decoder.load_state_dict(load_file(dec_path, device=str(device)), strict=True) + hm_decoder = hm_decoder.to(dtype) + + noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder='scheduler') + + sdpose_model = { + "unet": unet, + "vae": vae, + "empty_text_embed": empty_text_embed, + "decoder": hm_decoder, + "scheduler": noise_scheduler, + "keypoint_scheme": keypoint_scheme, + "device": device, + "unload_on_finish": unload_on_finish } + return (sdpose_model,) + + +# --- GroundingDINO Model Loader (Adapted from SAM2 Node) --- +# (global groundingdino_model_list definition) +groundingdino_model_dir_name = "grounding-dino" +groundingdino_model_list = { + "GroundingDINO_SwinT_OGC (694MB)": { + "config_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/GroundingDINO_SwinT_OGC.cfg.py", + "model_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/groundingdino_swint_ogc.pth", + }, + "GroundingDINO_SwinB (938MB)": { + "config_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/GroundingDINO_SwinB.cfg.py", + "model_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/groundingdino_swinb_cogcoor.pth", + }, +} if skeleton_type in ("Body", "OpenPose"): if not ULTRALYTICS_AVAILABLE: @@ -756,364 +972,440 @@ class IBBLoadPoseModel: return (model_info,) -class IBBRunPoseEstimation: - """ - IBB_POSE — Run Estimation + if dino_model_args.text_encoder_type == "bert-base-uncased": + dino_model_args.text_encoder_type = get_bert_base_uncased_model_path() - Detects persons (single or multi) and estimates poses. - Outputs: visualisation image AND/OR OpenPose JSON string. - """ + dino = local_groundingdino_build_model(dino_model_args) + checkpoint = torch.load( + get_local_filepath( + groundingdino_model_list[model_name]["model_url"], + groundingdino_model_dir_name, + ), map_location="cpu" # Load to CPU first + ) + dino.load_state_dict( + local_groundingdino_clean_state_dict(checkpoint["model"]), strict=False + ) + # Don't move to device here, let the Processor node handle it if needed + dino.eval() + return dino +# (modified for SDPose context) +class IBBGroundingDinoModelLoader: @classmethod def INPUT_TYPES(cls): return { "required": { - "ibb_pose_model": ("IBB_POSE_MODEL",), - "image": ("IMAGE",), - "detection_mode": (["single", "multi"],), - "output_type": (["PoseImage", "JSON", "Both"],), - "score_threshold": ("FLOAT", {"default": 0.3, "min": 0.05, "max": 0.95, "step": 0.05}), - "overlay_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), - "pose_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1}), - }, - "optional": { - "keep_face": ("BOOLEAN", {"default": True, "label_on": "Keep Face", "label_off": "No Face"}), - "keep_hands": ("BOOLEAN", {"default": True, "label_on": "Keep Hands", "label_off": "No Hands"}), - "keep_feet": ("BOOLEAN", {"default": True, "label_on": "Keep Feet", "label_off": "No Feet"}), - "save_json": ("BOOLEAN", {"default": False}), - "filename_prefix": ("STRING", {"default": "poses/ibb_pose"}), - }, - } - - RETURN_TYPES = ("IMAGE", "STRING") - RETURN_NAMES = ("pose_image", "pose_json") - FUNCTION = "estimate" - CATEGORY = "IBB_POSE" - OUTPUT_NODE = True - - @staticmethod - def _detect_yolo(yolo_det, img_bgr, conf=0.45): - H, W = img_bgr.shape[:2] - try: - results = yolo_det(img_bgr, classes=[0], verbose=False) - boxes = results[0].boxes - if boxes is not None and len(boxes) > 0: - bbs = boxes.xyxy.cpu().numpy().tolist() - confs = boxes.conf.cpu().numpy() - out = [b for b, c in zip(bbs, confs) if c >= conf] - if out: - return out - except Exception as exc: - print(f"IBB_POSE: YOLO detection error: {exc}") - return [[0, 0, W, H]] - - @staticmethod - def _run_yolo_pose(yolo_pose, img_bgr, detection_mode, yolo_det, threshold): - H, W = img_bgr.shape[:2] - - if detection_mode == "single": - try: - results = yolo_pose(img_bgr, verbose=False) - kps_data = results[0].keypoints - if kps_data is not None and len(kps_data) > 0: - kpts_arr = kps_data.xy.cpu().numpy() - conf_arr = kps_data.conf.cpu().numpy() - if len(kpts_arr) > 0: - best = int(np.argmax(conf_arr.sum(axis=1))) - return ([kpts_arr[best].astype(np.float32)], - [conf_arr[best].astype(np.float32)]) - except Exception as exc: - print(f"IBB_POSE: YOLO pose (single) error: {exc}") - return [], [] - - # multi - all_kpts, all_scores = [], [] - try: - results = yolo_pose(img_bgr, verbose=False) - kps_data = results[0].keypoints - if kps_data is not None and len(kps_data) > 0: - kpts_arr = kps_data.xy.cpu().numpy() - conf_arr = kps_data.conf.cpu().numpy() - for k, s in zip(kpts_arr, conf_arr): - all_kpts.append(k.astype(np.float32)) - all_scores.append(s.astype(np.float32)) - except Exception as exc: - print(f"IBB_POSE: YOLO pose (multi) error: {exc}") - return all_kpts, all_scores - - @staticmethod - def _run_wholebody(model_info, img_bgr, detection_mode, conf): - H, W = img_bgr.shape[:2] - pose_sess = model_info["pose_session"] - yolo_det = model_info.get("yolo_det") - det_sess = model_info.get("det_session") - - if detection_mode == "single": - kpts, scores = _run_dwpose_crop(pose_sess, img_bgr, None) - if kpts is None: - return [], [] - return [kpts], [scores] - - bboxes = [] - if yolo_det is not None: - try: - results = yolo_det(img_bgr, classes=[0], verbose=False) - boxes = results[0].boxes - if boxes is not None and len(boxes) > 0: - bbs = boxes.xyxy.cpu().numpy().tolist() - confs = boxes.conf.cpu().numpy() - bboxes = [b for b, c in zip(bbs, confs) if c >= conf] - except Exception as exc: - print(f"IBB_POSE: YOLO det for wholebody failed: {exc}") - - if not bboxes and det_sess is not None: - bboxes = _detect_persons_yolox(det_sess, img_bgr, conf) - - if not bboxes: - bboxes = [[0, 0, W, H]] - - all_kpts, all_scores = [], [] - for bb in bboxes: - k, s = _run_dwpose_crop(pose_sess, img_bgr, bb) - if k is not None: - all_kpts.append(k) - all_scores.append(s) - - return all_kpts, all_scores - - def estimate(self, ibb_pose_model, image, detection_mode, output_type, - score_threshold, overlay_alpha, pose_scale, - keep_face=True, keep_hands=True, keep_feet=True, - save_json=False, filename_prefix="poses/ibb_pose"): - - backend = ibb_pose_model.get("backend", "yolo_pose") - skel_type = ibb_pose_model.get("skeleton_type", "body") - threshold = score_threshold - - out_frames = [] - json_results = [] - - for b_idx in range(image.shape[0]): - frame_tensor = image[b_idx:b_idx+1] - img_bgr = _tensor_to_bgr(frame_tensor) - H, W = img_bgr.shape[:2] - - all_kpts: list = [] - all_scores: list = [] - - try: - if backend == "yolo_pose": - yolo_pose = ibb_pose_model["yolo_pose"] - yolo_det = ibb_pose_model.get("yolo_det") - all_kpts, all_scores = self._run_yolo_pose( - yolo_pose, img_bgr, detection_mode, yolo_det, threshold - ) - elif backend == "dwpose_onnx": - all_kpts, all_scores = self._run_wholebody( - ibb_pose_model, img_bgr, detection_mode, threshold - ) - else: - print(f"IBB_POSE: Unknown backend '{backend}'") - except Exception as exc: - print(f"IBB_POSE: Estimation error frame {b_idx}: {exc}") - traceback.print_exc() - - # Build pose image - if output_type in ("PoseImage", "Both"): - if overlay_alpha >= 1.0: - canvas = np.zeros_like(img_bgr) - else: - canvas = (img_bgr * (1.0 - overlay_alpha)).astype(np.uint8) - - if skel_type in ("body", "openpose"): - for kpts, scores in zip(all_kpts, all_scores): - if kpts.shape[0] == 17: - op_k, op_s = _coco17_to_openpose18(kpts, scores) - else: - op_k, op_s = kpts[:18], scores[:18] - canvas = _draw_body(canvas, op_k, op_s, threshold, pose_scale) - else: - for kpts, scores in zip(all_kpts, all_scores): - if kpts is not None: - canvas = _draw_wholebody( - canvas, kpts, scores, threshold, - keep_face, keep_hands, keep_feet, pose_scale - ) - - out_frames.append(_bgr_to_tensor(canvas)) - else: - out_frames.append(frame_tensor) - - # Build JSON - json_str = _build_openpose_json( - all_kpts, all_scores, W, H, - skeleton_type=skel_type, - threshold=threshold, - filter_low=True, - ) - json_results.append(json_str) - - # Save JSON - if save_json: - try: - if hasattr(folder_paths, "get_output_directory"): - base_out = folder_paths.get_output_directory() - else: - base_out = folder_paths.output_dir - prefix_dir = os.path.dirname(filename_prefix) - out_dir = os.path.join(base_out, prefix_dir) if prefix_dir else base_out - os.makedirs(out_dir, exist_ok=True) - base_name = os.path.basename(filename_prefix) - fname = f"{base_name}_{int(time.time())}.json" - fpath = os.path.join(out_dir, fname) - combined = ( - json_results[0] if len(json_results) == 1 - else json.dumps( - {"frames": [json.loads(j) for j in json_results]}, indent=2 - ) - ) - with open(fpath, "w", encoding="utf-8") as fh: - fh.write(combined) - print(f"IBB_POSE: JSON saved -> {fpath}") - except Exception as exc: - print(f"IBB_POSE: Could not save JSON: {exc}") - - pose_image_batch = torch.cat(out_frames, dim=0) - final_json = ( - json_results[-1] if json_results - else json.dumps({"people": [], "canvas_width": 0, "canvas_height": 0}) - ) - return (pose_image_batch, final_json) - - -class IBBPoseJsonToImage: - """ - IBB_POSE — JSON to Image - - Renders an OpenPose JSON string back to a pose visualisation image. - """ - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "pose_json": ("STRING", {"forceInput": True}), - "width": ("INT", {"default": 512, "min": 64, "max": 8192}), - "height": ("INT", {"default": 768, "min": 64, "max": 8192}), - "pose_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1}), - "threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.05}), - "keep_face": ("BOOLEAN", {"default": True}), - "keep_hands": ("BOOLEAN", {"default": True}), - "keep_feet": ("BOOLEAN", {"default": True}), + "model_name": (list_groundingdino_model(),), } } - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("pose_image",) - FUNCTION = "json_to_image" - CATEGORY = "IBB_POSE" + CATEGORY = "IBB_POSE" # Match SDPose category + FUNCTION = "load_groundingdino_model" + RETURN_TYPES = ("GROUNDING_DINO_MODEL",) - def json_to_image(self, pose_json, width, height, pose_scale, threshold, - keep_face, keep_hands, keep_feet): - canvas = np.zeros((height, width, 3), dtype=np.uint8) + def load_groundingdino_model(self, model_name): + dino_model = load_groundingdino_model(model_name) + # Wrap in a dictionary for potential future extensions + gd_model_wrapper = { + "model": dino_model, + "model_name": model_name, + } + return (gd_model_wrapper,) +class IBBPoseProcessor: + """ + ComfyUI node to run SDPose inference using true parallel batching. + Implements a 3-stage process: + 1. Collect: Detect all persons in all frames using YOLO. + 2. Batch Inference: Run SDPose model in batches on detected persons. + 3. Reconstruct: Draw poses back onto their original frames. + """ + + # 一个简单的数据类，用于跟踪检测到的人 + class DetectionJob: + """Track one detected person crop through preprocessing, inference, and restore.""" + + def __init__(self, frame_idx, person_idx, input_tensor, crop_info): + self.frame_idx = frame_idx + self.person_idx = person_idx + self.input_tensor = input_tensor + self.crop_info = crop_info + + # 这些字段将在推理后填充 + self.kpts = None + self.scores = None + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ibb_pose_model": ("IBB_POSE_MODEL",), + "images": ("IMAGE",), + "score_threshold": ("FLOAT", {"default": 0.3, "min": 0.1, "max": 0.9, "step": 0.05}), + "overlay_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}), + }, + "optional": { + "data_from_florence2": ("JSON",), + "grounding_dino_model": ("GROUNDING_DINO_MODEL",), + "prompt": ("STRING", {"default": "person .", "multiline": False}), + "gd_threshold": ("FLOAT", {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01}), + "yolo_model": ("YOLO_MODEL",), + "save_for_editor": ("BOOLEAN", {"default": False}), + "filename_prefix_edit": ("STRING", {"default": "poses/ibb_pose"}), + # --- 编辑选项 --- + "keep_face": ("BOOLEAN", {"default": True, "label_on": "Keep Face", "label_off": "Remove Face"}), + "keep_hands": ("BOOLEAN", {"default": True, "label_on": "Keep Hands", "label_off": "Remove Hands"}), + "keep_feet": ("BOOLEAN", {"default": True, "label_on": "Keep Feet", "label_off": "Remove Feet"}), + "scale_for_xinsr": ("BOOLEAN", {"default": False, "label_on": "Xinsr CN Scale", "label_off": "Default Scale"}), + # --- 核心修复：找回 Pose Scale 参数 --- + "pose_scale_factor": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1, "label": "Pose Size Scale"}), + "enable_confidence_filter": ("BOOLEAN", { + "default": True, + "label_on": "Filter Low Confidence", + "label_off": "Keep All Keypoints" + }), + } + + } + + RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT") + RETURN_NAMES = ("images", "pose_keypoint") + FUNCTION = "process_sequence" + CATEGORY = "IBB_POSE" + + def process_sequence( + self, + ibb_pose_model, + images, + score_threshold, + overlay_alpha, + batch_size=8, + grounding_dino_model=None, + prompt="person .", + gd_threshold=0.3, + yolo_model=None, + save_for_editor=False, + filename_prefix_edit="poses/ibb_pose", + data_from_florence2=None, + keep_face=True, + keep_hands=True, + keep_feet=True, + scale_for_xinsr=False, + pose_scale_factor=1.0, # <--- 必须包含此参数 + enable_confidence_filter=True + ): + + import gc + + device = ibb_pose_model["device"] + keypoint_scheme = ibb_pose_model["keypoint_scheme"] + input_size = (768, 1024) + + B, H, W, C = images.shape + images_np = images.cpu().numpy() + print(f"IBB_POSE: Received {B} frames. Starting 3-stage process with batch_size={batch_size}.") + + # 检查模型设备 try: - data = json.loads(pose_json) - people = data.get("people", []) - cw = data.get("canvas_width", width) - ch = data.get("canvas_height", height) - scale_x = width / max(cw, 1) - scale_y = height / max(ch, 1) + for key in ["unet", "vae", "decoder"]: + if key in ibb_pose_model: + current_device = next(ibb_pose_model[key].parameters()).device + if current_device != device: + ibb_pose_model[key].to(device) + except Exception as e: + print(f"IBB_POSE: CRITICAL ERROR moving models to device: {e}. Aborting.") + raise e + + class MockPipeline: + def __init__(self, model_dict): + self.unet = model_dict["unet"] + self.vae = model_dict["vae"] + self.decoder = model_dict["decoder"] + self.scheduler = model_dict["scheduler"] + self.empty_text_embed = model_dict["empty_text_embed"] + self.device = model_dict["device"] - for person in people: - raw = person.get("pose_keypoints_2d", []) - if raw: - n_k = len(raw) // 3 - op_k = np.zeros((n_k, 2), dtype=np.float32) - op_s = np.zeros(n_k, dtype=np.float32) - for i in range(n_k): - op_k[i, 0] = raw[i * 3] * scale_x - op_k[i, 1] = raw[i * 3 + 1] * scale_y - op_s[i] = raw[i * 3 + 2] - canvas = _draw_body(canvas, op_k, op_s, threshold, pose_scale) + @torch.no_grad() + def __call__(self, rgb_in, timesteps, test_cfg, **kwargs): + unet_dtype = self.unet.dtype + bsz = rgb_in.shape[0] + + rgb_latent = self.vae.encode(rgb_in).latent_dist.sample() * 0.18215 + rgb_latent = rgb_latent.to(dtype=unet_dtype) + t = torch.tensor(timesteps, device=self.device).long() + text_embed = self.empty_text_embed.repeat((bsz, 1, 1)).to(dtype=unet_dtype) + task_emb_anno = torch.tensor([1, 0]).float().unsqueeze(0).to(self.device) + task_emb_anno = torch.cat([torch.sin(task_emb_anno), torch.cos(task_emb_anno)], dim=-1).repeat(bsz, 1) + task_emb_anno = task_emb_anno.to(dtype=unet_dtype) + + feat = self.unet(rgb_latent, t, text_embed, class_labels=task_emb_anno, return_dict=False, return_decoder_feats=True) + return self.decoder.predict((feat,), None, test_cfg=test_cfg) - if keep_face: - raw_f = person.get("face_keypoints_2d", []) - if raw_f: - n_f = len(raw_f) // 3 - f_r = max(1, int(3 * pose_scale)) - for i in range(n_f): - s = raw_f[i * 3 + 2] - if s < threshold: - continue - x = int(raw_f[i * 3] * scale_x) - y = int(raw_f[i * 3 + 1] * scale_y) - if 0 < x < width and 0 < y < height: - cv2.circle(canvas, (x, y), f_r, (255, 255, 255), -1) + pipeline = MockPipeline(ibb_pose_model) + + # --- 步骤 1: 收集 --- + all_jobs = [] + + # Florence2 logic + input_bboxes_per_frame_f2 = [None] * B + used_florence2 = False + if data_from_florence2 is not None: + if isinstance(data_from_florence2, list) and len(data_from_florence2) == B: + valid_f2_data = True + parsed_bboxes_list = [] + for i, frame_data in enumerate(data_from_florence2): + if isinstance(frame_data, dict) and 'bboxes' in frame_data and isinstance(frame_data['bboxes'], list): + bboxes_for_frame = [list(map(float, box)) for box in frame_data['bboxes'] if len(box) == 4] + parsed_bboxes_list.append(bboxes_for_frame) + else: + if frame_data is None or (isinstance(frame_data, dict) and not frame_data.get('bboxes')): + parsed_bboxes_list.append([]) + else: + valid_f2_data = False + break + if valid_f2_data: + input_bboxes_per_frame_f2 = parsed_bboxes_list + used_florence2 = True - if keep_hands: - f_lw = max(1, int(2 * pose_scale)) - f_r = max(1, int(3 * pose_scale)) - for hand_key in ("hand_left_keypoints_2d", "hand_right_keypoints_2d"): - raw_h = person.get(hand_key, []) - if not raw_h: - continue - n_h = len(raw_h) // 3 - hk = np.zeros((n_h, 2), dtype=np.float32) - hs = np.zeros(n_h, dtype=np.float32) - for i in range(n_h): - hk[i, 0] = raw_h[i * 3] * scale_x - hk[i, 1] = raw_h[i * 3 + 1] * scale_y - hs[i] = raw_h[i * 3 + 2] - for ie, (ea, eb) in enumerate(HAND_EDGES): - if ea >= n_h or eb >= n_h: - continue - if hs[ea] < threshold or hs[eb] < threshold: - continue - x1, y1 = int(hk[ea, 0]), int(hk[ea, 1]) - x2, y2 = int(hk[eb, 0]), int(hk[eb, 1]) - if (0 < x1 < width and 0 < y1 < height - and 0 < x2 < width and 0 < y2 < height): - red, green, blue = colorsys.hsv_to_rgb(ie / len(HAND_EDGES), 1.0, 1.0) - cv2.line(canvas, (x1, y1), (x2, y2), - (int(blue * 255), int(green * 255), int(red * 255)), f_lw) - for i in range(n_h): - if hs[i] < threshold: - continue - x, y = int(hk[i, 0]), int(hk[i, 1]) - if 0 < x < width and 0 < y < height: - cv2.circle(canvas, (x, y), f_r, (0, 0, 255), -1) + for frame_idx in range(B): + original_image_rgb = (images_np[frame_idx] * 255).astype(np.uint8) + original_image_bgr = cv2.cvtColor(original_image_rgb, cv2.COLOR_RGB2BGR) + H_curr, W_curr = original_image_rgb.shape[:2] - if keep_feet: - raw_ft = person.get("foot_keypoints_2d", []) - if raw_ft: - n_ft = len(raw_ft) // 3 - cr = max(1, int(4 * pose_scale)) - for i in range(n_ft): - s = raw_ft[i * 3 + 2] - if s < threshold: - continue - x = int(raw_ft[i * 3] * scale_x) - y = int(raw_ft[i * 3 + 1] * scale_y) - if 0 <= x < width and 0 <= y < height: - cv2.circle(canvas, (x, y), cr, - BODY_COLORS[i % len(BODY_COLORS)], -1) + bboxes = [] + detection_source = "None" - except Exception as exc: - print(f"IBB_POSE: json_to_image error: {exc}") - traceback.print_exc() + if used_florence2 and input_bboxes_per_frame_f2[frame_idx] is not None: + bboxes = input_bboxes_per_frame_f2[frame_idx] + detection_source = "Florence2" + + if not bboxes and grounding_dino_model is not None: + if prompt and prompt.strip(): + try: + pil_image = Image.fromarray(original_image_rgb) + gd_bboxes = groundingdino_predict(grounding_dino_model, pil_image, prompt, gd_threshold) + if gd_bboxes: + bboxes = gd_bboxes + detection_source = "GroundingDINO" + except Exception: pass - return (_bgr_to_tensor(canvas),) + if not bboxes and yolo_model is not None and YOLO_AVAILABLE: + try: + results = yolo_model(original_image_bgr, verbose=False) + yolo_bboxes = [] + for result in results: + if result.boxes is not None: + for box in result.boxes: + if int(box.cls[0]) == 0 and float(box.conf[0]) > 0.5: + yolo_bboxes.append(box.xyxy[0].cpu().numpy().tolist()) + if yolo_bboxes: + bboxes = yolo_bboxes + detection_source = "YOLO" + except Exception: pass + if not bboxes: + bboxes = [[0, 0, W_curr, H_curr]] + detection_source = "Full Image" + for person_idx, bbox in enumerate(bboxes): + if not (isinstance(bbox, (list, tuple)) and len(bbox) == 4): continue + x1, y1, x2, y2 = bbox + if not (x2 > x1 and y2 > y1): continue + input_tensor, _, crop_info = preprocess_image_for_sdpose(original_image_bgr, bbox, input_size) + all_jobs.append(self.DetectionJob(frame_idx, person_idx, input_tensor, crop_info)) + + total_detections = len(all_jobs) + print(f"IBB_POSE: Stage 1 complete. Found {total_detections} persons.") + + # --- 步骤 2: 并行推理 (批量处理) - 内存优化版 --- + model_management.soft_empty_cache() + + for i in range(0, total_detections, batch_size): + batch_jobs = all_jobs[i : i + batch_size] + current_batch_size = len(batch_jobs) + + batch_tensors = torch.cat([job.input_tensor for job in batch_jobs], dim=0).to(device) + + with torch.no_grad(): + out = pipeline(batch_tensors, timesteps=[999], test_cfg={'flip_test': False}, show_progress_bar=False, mode="inference") + + for j in range(current_batch_size): + kpts_raw = out[j].keypoints[0] + if hasattr(kpts_raw, 'cpu'): + batch_jobs[j].kpts = kpts_raw.cpu().numpy() + else: + batch_jobs[j].kpts = kpts_raw + + scores_raw = out[j].keypoint_scores[0] + if hasattr(scores_raw, 'cpu'): + batch_jobs[j].scores = scores_raw.cpu().numpy() + else: + batch_jobs[j].scores = scores_raw + + batch_jobs[j].input_tensor = None + + del batch_tensors + del out + + model_management.soft_empty_cache() + + # --- 步骤 3: 重组 (绘图和JSON) --- + print("IBB_POSE: Stage 3/3 - Reconstructing frames...") + + frame_data = [] + for i in range(B): + frame_data.append({ + "canvas": np.zeros_like((images_np[i] * 255).astype(np.uint8), order='C'), + "all_keypoints": [], + "all_scores": [], + }) + + for job in all_jobs: + frame_idx = job.frame_idx + + kpts_original = restore_keypoints_to_original(job.kpts, job.crop_info, input_size, (W, H)) + scores = job.scores + + if keypoint_scheme == "body": + kpts_final = kpts_original + scores_final = scores + frame_data[frame_idx]["canvas"] = draw_body17_keypoints_openpose_style( + frame_data[frame_idx]["canvas"], kpts_final, scores_final, + threshold=score_threshold, + scale_for_xinsr=scale_for_xinsr, + pose_scale=pose_scale_factor # <--- 传递参数 + ) + + else: # wholebody + kpts_final = kpts_original.copy() + scores_final = scores.copy() + if len(kpts_original) >= 17: + neck = (kpts_original[5] + kpts_original[6]) / 2 + neck_score = min(scores[5], scores[6]) if scores[5] > 0.3 and scores[6] > 0.3 else 0 + kpts_final = np.insert(kpts_original, 17, neck, axis=0) + scores_final = np.insert(scores, 17, neck_score) + + mmpose_idx = np.array([17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3]) + openpose_idx = np.array([1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17]) + temp_kpts = kpts_final.copy(); temp_scores = scores_final.copy() + temp_kpts[openpose_idx] = kpts_final[mmpose_idx] + temp_scores[openpose_idx] = scores_final[mmpose_idx] + kpts_final = temp_kpts; scores_final = temp_scores + + if len(kpts_final) >= 134: + if not keep_face: + kpts_final[24:92] = 0.0; scores_final[24:92] = 0.0 + if not keep_hands: + kpts_final[92:134] = 0.0; scores_final[92:134] = 0.0 + if not keep_feet: + kpts_final[18:24] = 0.0; scores_final[18:24] = 0.0 + + frame_data[frame_idx]["canvas"] = draw_wholebody_keypoints_openpose_style( + frame_data[frame_idx]["canvas"], kpts_final, scores_final, + threshold=score_threshold, + scale_for_xinsr=scale_for_xinsr, + pose_scale=pose_scale_factor # <--- 传递参数 + ) + + frame_data[frame_idx]["all_keypoints"].append(kpts_final) + frame_data[frame_idx]["all_scores"].append(scores_final) + + result_images = [] + all_frames_json_data = [] + + for frame_idx in range(B): + original_image_rgb = (images_np[frame_idx] * 255).astype(np.uint8) + pose_canvas = frame_data[frame_idx]["canvas"] + result_image = cv2.addWeighted(original_image_rgb, 1.0 - overlay_alpha, pose_canvas, overlay_alpha, 0) + result_images.append(result_image) + + frame_json = convert_to_openpose_json( + frame_data[frame_idx]["all_keypoints"], + frame_data[frame_idx]["all_scores"], + W, H, keypoint_scheme + ) + all_frames_json_data.append(frame_json) + + # --- 修复后的保存编辑器 JSON 逻辑 --- + if save_for_editor: + try: + data_to_save = convert_to_loader_json( + frame_data[frame_idx]["all_keypoints"], + frame_data[frame_idx]["all_scores"], + W, H, keypoint_scheme, score_threshold, enable_confidence_filter + ) + + output_dir = folder_paths.get_output_directory() + is_batch = B > 1 + + if is_batch: + filename_to_use = f"{filename_prefix_edit}_frame{frame_idx:06d}" + full_output_folder, base_filename, _, _, _ = folder_paths.get_save_image_path(filename_to_use, output_dir, W, H) + os.makedirs(full_output_folder, exist_ok=True) + final_filename = f"{base_filename}.json" + file_path = os.path.join(full_output_folder, final_filename) + + else: + full_output_folder, base_filename, _, subfolder, _ = folder_paths.get_save_image_path(filename_prefix_edit, output_dir, W, H) + + os.makedirs(full_output_folder, exist_ok=True) + + counter = 1 + try: + existing_files = [f for f in os.listdir(full_output_folder) if f.startswith(base_filename + "_") and f.endswith(".json")] + if existing_files: + max_counter = 0 + for f in existing_files: + try: + num_str = f[len(base_filename)+1:-5] + num = int(num_str) + if num > max_counter: + max_counter = num + except ValueError: + continue + counter = max_counter + 1 + except FileNotFoundError: + pass + + final_filename = f"{base_filename}_{counter:05d}.json" + file_path = os.path.join(full_output_folder, final_filename) + + with open(file_path, 'w') as f: + json.dump(data_to_save, f, indent=4) + + if not is_batch: + print(f"IBB_POSE: Saved pose JSON to: {file_path}") + + except Exception as e: + print(f"IBB_POSE: ERROR Failed to save editor JSON for frame {frame_idx}: {e}") + + result_tensor = torch.from_numpy(np.stack(result_images, axis=0).astype(np.float32) / 255.0) + + # --- 最终清理 --- + if ibb_pose_model.get("unload_on_finish", False): + offload_device = torch.device("cpu") + for key in ["unet", "vae", "decoder"]: + if key in ibb_pose_model: ibb_pose_model[key].to(offload_device) + + del all_jobs + del pipeline + gc.collect() + model_management.soft_empty_cache() + + return (result_tensor, all_frames_json_data) + +# --- Node Mappings for ComfyUI --- NODE_CLASS_MAPPINGS = { - "IBBLoadPoseModel": IBBLoadPoseModel, - "IBBRunPoseEstimation": IBBRunPoseEstimation, - "IBBPoseJsonToImage": IBBPoseJsonToImage, + "IBBPoseModelLoader": IBBPoseModelLoader, + "IBBPoseProcessor": IBBPoseProcessor, + "IBBYOLOModelLoader": IBBYOLOModelLoader, + "IBBGroundingDinoModelLoader": IBBGroundingDinoModelLoader, } NODE_DISPLAY_NAME_MAPPINGS = { - "IBBLoadPoseModel": "IBB Pose — Load Model", - "IBBRunPoseEstimation": "IBB Pose — Run Estimation", - "IBBPoseJsonToImage": "IBB Pose — JSON to Image", + "IBBPoseModelLoader": "IBB Pose — Load Model", + "IBBPoseProcessor": "IBB Pose — Run SDPose Estimation", + "IBBYOLOModelLoader": "IBB Pose — Load YOLO Model", + "IBBGroundingDinoModelLoader": "IBB Pose — Load GroundingDINO Model", } diff --git a/models/HeatmapHead.py b/models/HeatmapHead.py new file mode 100644 index 0000000..b5268b7 --- /dev/null +++ b/models/HeatmapHead.py @@ -0,0 +1,64 @@ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk +# Date: Jun. 2025 +# Description: SDPose heatmap head factory using mmpose `HeatmapHead` + `UDPHeatmap` codec. +# - mode="body" → 17 keypoints, in_channels=320 +# - mode="wholebody"→ 133 keypoints, in_channels=640 +# Defaults: image_size=(768,1024), scale=4 (→ heatmap size = (w/4, h/4)), sigma=6. +# License: MIT License (see LICENSE for details) + +from mmpose.registry import MODELS + +def get_heatmap_head(mode = "body"): + + if mode == "body": + + image_size = [768, 1024] ## width x height + sigma = 6 ## sigma is 2 for 256 + scale = 4 + + embed_dim = 320 + num_keypoints = 17 + + codec = dict( + type='UDPHeatmap', input_size=(image_size[0], image_size[1]), heatmap_size=(int(image_size[0]/scale), int(image_size[1]/scale)), sigma=sigma) ## sigma is 2 for 256 + + head_cfg = dict( + type='HeatmapHead', + in_channels=embed_dim, + out_channels=num_keypoints, + deconv_out_channels=(320,), + deconv_kernel_sizes=(4,), + conv_out_channels=(320,), + conv_kernel_sizes=(1,), + loss=dict(type='KeypointMSELoss', use_target_weight=True), + decoder=codec + ) + + elif mode == "wholebody": + + image_size = [768, 1024] ## width x height + sigma = 6 ## sigma is 2 for 256 + scale = 4 + + embed_dim = 640 + num_keypoints = 133 + + codec = dict( + type='UDPHeatmap', input_size=(image_size[0], image_size[1]), heatmap_size=(int(image_size[0]/scale), int(image_size[1]/scale)), sigma=sigma) ## sigma is 2 for 256 + + head_cfg = dict( + type='HeatmapHead', + in_channels=embed_dim, + out_channels=num_keypoints, + deconv_out_channels=(640,), + deconv_kernel_sizes=(4,), + conv_out_channels=(640,), + conv_kernel_sizes=(1,), + loss=dict(type='KeypointMSELoss', use_target_weight=True), + decoder=codec + ) + + head = MODELS.build(head_cfg) + + return head \ No newline at end of file diff --git a/models/ModifiedUNet.py b/models/ModifiedUNet.py new file mode 100644 index 0000000..f785a63 --- /dev/null +++ b/models/ModifiedUNet.py @@ -0,0 +1,69 @@ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com +# Date: Feb. 2025 + +""" + +Description: + + SDPose UNet Forward Modifier + + Temporarily overrides `UNet.forward` to expose decoder features via forward hooks. + - keypoint_scheme: "body" (COCO-17) or "wholebody" (COCO-WholeBody-133) + - return_decoder_feats: if True, returns a selected decoder feature (Tensor); + otherwise returns the original UNet output. +""" + +def new_forward_kpt17(self, *args, return_decoder_feats=False, **kwargs): + self._decoder_feats = [] + + def hook_fn(module, input, output): + self._decoder_feats.append(output) + + handles = [blk.register_forward_hook(hook_fn) for blk in self.up_blocks] + + out = self._old_forward(*args, **kwargs) + + for h in handles: + h.remove() + + if return_decoder_feats: + + feats = self._decoder_feats[::-1] + return feats[0] + else: + return out + +def new_forward_kpt133(self, *args, return_decoder_feats=False, **kwargs): + self._decoder_feats = [] + + def hook_fn(module, input, output): + self._decoder_feats.append(output) + + handles = [blk.register_forward_hook(hook_fn) for blk in self.up_blocks] + + out = self._old_forward(*args, **kwargs) + + for h in handles: + h.remove() + + if return_decoder_feats: + + feats = self._decoder_feats[::-1] + return feats[1] + else: + return out + +def Modified_forward(unet, keypoint_scheme = "body"): + + if keypoint_scheme == "body": + + unet._old_forward = unet.forward + unet.forward = new_forward_kpt17.__get__(unet, unet.__class__) + + elif keypoint_scheme == "wholebody": + + unet._old_forward = unet.forward + unet.forward = new_forward_kpt133.__get__(unet, unet.__class__) + + return unet \ No newline at end of file diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..46f419a --- /dev/null +++ b/models/__init__.py @@ -0,0 +1 @@ +# Models package for SDPose OOD \ No newline at end of file diff --git a/pipelines/SDPose_D_Pipeline.py b/pipelines/SDPose_D_Pipeline.py new file mode 100644 index 0000000..f1d05bc --- /dev/null +++ b/pipelines/SDPose_D_Pipeline.py @@ -0,0 +1,404 @@ +# Author T. S. Liang @ Rama Alpaca, July. 2024. +# Email: tsliang2001@gmail.com & shuangliang@ramaalpaca.com +# Adapted from Marigold: https://github.com/prs-eth/marigold and MMPose: https://github.com/open-mmlab/mmpose +# License: MIT License (see LICENSE for details) + +""" + Description: + + This script is the pipeline utilizing Stable Diffusion Model for the pose perception. + Adaptations have been made to ensure full compatibility with MMPose-style + decoder heads, enabling seamless integration with existing heatmap-based + keypoint decoders and evaluation protocols. +""" + +from typing import Union, Optional + +import torch +import torch.nn as nn + +from diffusers import ( + AutoencoderKL, + DiffusionPipeline, + UNet2DConditionModel, + DDIMScheduler, + DDPMScheduler, +) + +from PIL import Image + +from tqdm.auto import tqdm +from transformers import CLIPTextModel, CLIPTokenizer + +from itertools import zip_longest +from mmpose.utils.typing import (ConfigType, InstanceList, OptConfigType, + OptMultiConfig, PixelDataList, SampleList) + +from mmpose.models.heads.heatmap_heads.heatmap_head import HeatmapHead + +class SDPose_D_Pipeline(DiffusionPipeline): + + """ + A diffusion-based human pose estimation pipeline that reconstructs pose heatmaps + from RGB images using a Stable-Diffusion-style latent denoising process. + + NOTE: + This pipeline assumes the scheduler/unet are configured for x0 prediction, + i.e., `scheduler.config.prediction_type == "sample"`. In this setting the UNet + directly predicts the clean latent x0 at a fixed diffusion step. + + Args: + unet (`UNet2DConditionModel`): + Conditional UNet used to predict the clean latent `x0`. + Inputs: noisy pose/image latents and conditioning (image latent, text embeds). + vae (`AutoencoderKL`): + Variational Autoencoder for encoding RGB images to latents and for decoding + latents (image/heatmap) back to pixel space as needed. + text_encoder (`CLIPTextModel`): + Text encoder to obtain conditioning embeddings. For text-free pose tasks, + empty prompts are typically used. + tokenizer (`CLIPTokenizer`): + Tokenizer paired with the text encoder (tokenizes empty/neutral prompts). + decoder (`HeatmapHead`): + A lightweight head mapping denoised latent features to pose heatmaps, + typically with shape `[B, num_keypoints, H, W]`. + scheduler (`Union[DDIMScheduler, DDPMScheduler]`): + Diffusion scheduler defining the diffusion process configuration + (only used to set the fixed timestep). + Its `prediction_type` **must be `'sample'`** to ensure x₀ prediction. + + Returns: + `SDPose_D_Pipeline`: + Initialized pipeline ready for inference/training. + """ + + def __init__( + self, + unet: UNet2DConditionModel, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + decoder: HeatmapHead, + scheduler: Union[DDIMScheduler, DDPMScheduler] + ): + + super().__init__() + + self.register_modules( + unet = unet, + vae = vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + decoder = decoder, + scheduler = scheduler + ) + + self.rgb_latent_scale_factor = 0.18215 + + self.test_cfg = { + 'output_keypoint_indices': None, + 'flip_test': True, + 'flip_mode': 'heatmap', + 'shift_heatmap': False + } + + @torch.no_grad() + def __call__( + self, + rgb_in: Union[Image.Image, torch.Tensor], + show_progress_bar: bool = True, + timesteps = None, + mode = None, + data_samples = None, + test_cfg = None + ): + + """ + Run a forward inference through the SDPose-D pipeline. + + This function performs **x₀ prediction** using a Stable-Diffusion-style UNet + conditioned on the RGB image latent, and optionally returns decoded pose heatmaps + or full structured pose predictions (MMPose-compatible). + + Args: + rgb_in (`PIL.Image.Image` or `torch.Tensor`): + Input RGB image or batch tensor of shape `[B, 3, H, W]`. + If a PIL image is provided, it should be preprocessed to a normalized tensor + compatible with the VAE encoder input range. + show_progress_bar (`bool`, optional, defaults to `True`): + Whether to display a progress bar during diffusion inference. + timesteps (`int` or `list[int]`, optional): + Diffusion timestep(s) to use. + mode (`str`, optional): + - `"predict"`: Use the decoder's `predict()` method to generate keypoint + predictions and attach them to `data_samples` (for evaluation). + - `None`: Decode the denoised latent directly to a heatmap tensor `[B, K, H, W]` + without MMPose post-processing. + data_samples (`SampleList`, optional): + A list of MMPose `PoseDataSample` objects required when `mode="predict"`. + test_cfg (`dict`, optional): + Configuration overrides for test-time behavior. + Examples: `{'flip_test': True, 'shift_heatmap': False}`. + + Returns: + - If `mode == "predict"`: + `SampleList`: + The input `data_samples` with updated `pred_instances` and optional `pred_fields` + containing the decoded keypoints and heatmaps. + - Else: + `torch.Tensor`: + The predicted pose heatmap tensor of shape `[B, num_keypoints, H, W]`. + + Notes: + - This pipeline operates in **x₀ prediction mode**, where the UNet directly predicts + the clean latent representation instead of ε or v. + - The pose diffusion process uses a **single, fixed timestep** (default t=999), + meaning there is no iterative denoising loop. + - `test_cfg` supports `flip_test` inference (horizontal ensemble of original and flipped images). + - `scheduler.config.prediction_type` must be set to `"sample"` for correct operation. + """ + + # ------------- Predicting pose --------------- + + # Predict pose images + + self.encode_empty_text() + + # Update test_cfg if provided + if test_cfg is not None: + self.test_cfg.update(test_cfg) + + heatmap_pred = self.single_infer( + rgb_in = rgb_in, + show_pbar = show_progress_bar, + timesteps = timesteps, + mode = mode, + data_samples = data_samples + ) + + return heatmap_pred + + @torch.no_grad() + def single_infer( + self, + rgb_in: torch.Tensor, + timesteps, + show_pbar: bool, + mode = None, + data_samples = None + )-> torch.Tensor: + + device = self.unet.device + rgb_in = rgb_in.to(device) + bsz = rgb_in.shape[0] + + self.scheduler.set_timesteps(timesteps = timesteps, device=device) + + timesteps = torch.tensor(timesteps, device=device).long() + + # STEP 1: Encode image + if self.test_cfg.get('flip_test', False): + # Flip test: process both original and flipped images + rgb_in_flipped = torch.flip(rgb_in, dims=[3]) # Flip horizontally (W dimension) + + # Encode both original and flipped images + rgb_latent = self.encode_rgb(rgb_in) + rgb_latent_flipped = self.encode_rgb(rgb_in_flipped) + + # Process both latents + feats = [] + for rgb_latent_curr in [rgb_latent, rgb_latent_flipped]: + # Denoising loop for current latent + if show_pbar: + iterable = tqdm( + enumerate(timesteps), + total=len(timesteps), + leave=False, + desc=" " * 4 + "Diffusion denoising", + ) + else: + iterable = enumerate(timesteps) + + self.encode_empty_text() + text_embed = self.empty_text_embed.repeat((bsz, 1, 1)) + + for i, t in iterable: + task_emb_anno = torch.tensor([1, 0]).float().unsqueeze(0).to(device) + task_emb_anno = torch.cat([torch.sin(task_emb_anno), torch.cos(task_emb_anno)], dim=-1).repeat(bsz, 1) + + if self.scheduler.config.prediction_type == "sample": + feat = self.unet( + rgb_latent_curr, t, text_embed, class_labels=task_emb_anno, return_dict=False, return_decoder_feats=True + ) + else: + raise + + feats.append((feat,)) # Wrap each feat in a tuple to match expected format + + else: + # Original single image processing + rgb_latent = self.encode_rgb(rgb_in) + + # Denoising loop + if show_pbar: + iterable = tqdm( + enumerate(timesteps), + total=len(timesteps), + leave=False, + desc=" " * 4 + "Diffusion denoising", + ) + else: + iterable = enumerate(timesteps) + + self.encode_empty_text() + text_embed = self.empty_text_embed.repeat((bsz, 1, 1)) + + for i, t in iterable: + task_emb_anno = torch.tensor([1, 0]).float().unsqueeze(0).to(device) + task_emb_anno = torch.cat([torch.sin(task_emb_anno), torch.cos(task_emb_anno)], dim=-1).repeat(bsz, 1) + + if self.scheduler.config.prediction_type == "sample": + feat = self.unet( + rgb_latent, t, text_embed, class_labels=task_emb_anno, return_dict=False, return_decoder_feats=True + ) + else: + raise + + feats = (feat,) # Wrap single tensor in tuple to match expected format for prediction mode + feat_for_training = feat # Keep original for training mode + + if mode == "predict": + + if isinstance(self.decoder, (nn.parallel.DistributedDataParallel, )): + + preds = self.decoder.module.predict(feats, data_samples, test_cfg=self.test_cfg) + + else: + + preds = self.decoder.predict(feats, data_samples, test_cfg=self.test_cfg) + + if isinstance(preds, tuple): + batch_pred_instances, batch_pred_fields = preds + else: + batch_pred_instances = preds + batch_pred_fields = None + + return self.add_pred_to_datasample(batch_pred_instances=batch_pred_instances, batch_data_samples=data_samples, batch_pred_fields=batch_pred_fields) + + elif mode == "inference": + + preds = self.decoder.predict(feats, data_samples, test_cfg=self.test_cfg) + + return preds + + else: + if self.test_cfg.get('flip_test', False): + + if isinstance(feats, list): + # feats is [(feat_orig,), (feat_flipped,)], extract the first one for training + heatmap_pred = self.decoder(feats[0][0]) # Use feat_orig for training + else: + # feats is a single tuple (feat,), extract the tensor + heatmap_pred = self.decoder(feats[0]) + else: + # In training mode, decoder expects raw features, not wrapped in tuples + heatmap_pred = self.decoder(feat_for_training) # Use the original feat tensor + + return heatmap_pred + + def encode_empty_text(self): + + """ + Encode text embedding for empty prompt + """ + + prompt = "" + text_inputs = self.tokenizer( + prompt, + padding="do_not_pad", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(self.text_encoder.device) + self.empty_text_embed = self.text_encoder(text_input_ids)[0] + + def encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor: + """ + encode RGB image into latent. + + Args: + + rgb_in (`torch.Tensor`): + Input RGB image to be encoded. + + Returns: + `torch.Tensor`: Image latent. + """ + + rgb_latent = self.vae.encode(rgb_in).latent_dist.sample() + rgb_latent = rgb_latent * self.rgb_latent_scale_factor + + return rgb_latent + + def add_pred_to_datasample(self, batch_pred_instances: InstanceList, + batch_pred_fields: Optional[PixelDataList], + batch_data_samples: SampleList) -> SampleList: + + # from https://github.com/open-mmlab/mmpose/blob/759b39c13fea6ba094afc1fa932f51dc1b11cbf9/mmpose/models/pose_estimators/topdown.py#L122 + + """Add predictions into data samples. + + Args: + batch_pred_instances (List[InstanceData]): The predicted instances + of the input data batch + batch_pred_fields (List[PixelData], optional): The predicted + fields (e.g. heatmaps) of the input batch + batch_data_samples (List[PoseDataSample]): The input data batch + + Returns: + List[PoseDataSample]: A list of data samples where the predictions + are stored in the ``pred_instances`` field of each data sample. + """ + assert len(batch_pred_instances) == len(batch_data_samples) + if batch_pred_fields is None: + batch_pred_fields = [] + output_keypoint_indices = self.test_cfg.get('output_keypoint_indices', + None) + + for pred_instances, pred_fields, data_sample in zip_longest( + batch_pred_instances, batch_pred_fields, batch_data_samples): + + gt_instances = data_sample.gt_instances + + bbox_centers = gt_instances.bbox_centers + bbox_scales = gt_instances.bbox_scales + input_size = data_sample.metainfo['input_size'] + + pred_instances.keypoints = pred_instances.keypoints / input_size \ + * bbox_scales + bbox_centers - 0.5 * bbox_scales + + if output_keypoint_indices is not None: + + num_keypoints = pred_instances.keypoints.shape[1] + for key, value in pred_instances.all_items(): + if key.startswith('keypoint'): + pred_instances.set_field( + value[:, output_keypoint_indices], key) + + pred_instances.bboxes = gt_instances.bboxes + pred_instances.bbox_scores = gt_instances.bbox_scores + + data_sample.pred_instances = pred_instances + + if pred_fields is not None: + if output_keypoint_indices is not None: + + for key, value in pred_fields.all_items(): + if value.shape[0] != num_keypoints: + continue + pred_fields.set_field(value[output_keypoint_indices], + key) + data_sample.pred_fields = pred_fields + + return batch_data_samples diff --git a/pipelines/__init__.py b/pipelines/__init__.py new file mode 100644 index 0000000..3fd4720 --- /dev/null +++ b/pipelines/__init__.py @@ -0,0 +1 @@ +# Pipelines package for SDPose OOD \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index ec4a97c..783ae1b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ # Core (required) opencv-python>=4.8.0 numpy>=1.24.0 +opencv-python>=4.8.0 Pillow>=10.0.0 # Optional backend packages (install manually when needed) diff --git a/scripts/eval.sh b/scripts/eval.sh new file mode 100644 index 0000000..d73d2dd --- /dev/null +++ b/scripts/eval.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# ============================================ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk +# Date: Aug. 2025 +# Description: Launch script for evaluating SDPose on pose estimation benchmarks. +# License: MIT License (see LICENSE for details) +# ============================================ + +LAUNCH_EVAL() { + + cd .. + + # ------------------ Configurable Parameters ------------------ + dataset_name='COCO' + keypoint_scheme='body' # or 'wholebody' + dataset_root='/data/coding/datasets' + ann_file='/data/coding/datasets/HumanArt/annotations/validation_humanart.json' + checkpoint_path='/data/coding/outputs/COCO_813' + eval_batch_size=16 + dataloader_num_workers=16 + timestep=999 + + # ------------------ Runtime Environment ------------------ + CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --mixed_precision="bf16"\ + --multi_gpu \ + eval/launch_evaluator.py \ + --dataset_name $dataset_name \ + --keypoint_scheme $keypoint_scheme \ + --dataset_root $dataset_root \ + --checkpoint_path $checkpoint_path \ + --ann_file $ann_file \ + --eval_batch_size $eval_batch_size \ + --dataloader_num_workers $dataloader_num_workers \ + --enable_xformers_memory_efficient_attention \ + --timestep $timestep +} + +# ------------- Entry ------------- +LAUNCH_EVAL diff --git a/utils/EvalMetric.py b/utils/EvalMetric.py new file mode 100644 index 0000000..fb780fc --- /dev/null +++ b/utils/EvalMetric.py @@ -0,0 +1,24 @@ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com +# Date: Feb. 2025 +# Description: Build COCO-style keypoint metrics for SDPose evaluation. +# License: MIT License (see LICENSE for details) + + +from mmpose.evaluation.metrics import CocoMetric, CocoWholeBodyMetric + +def get_metric(ann_file, mode = "body"): + + if mode == "body": + + metric = CocoMetric(iou_type='keypoints', ann_file = ann_file) + + elif mode == "wholebody": + + metric = CocoWholeBodyMetric(iou_type='keypoints', ann_file = ann_file) + + else: + + print("Error, please specify the correct type of keypoint scheme.") + + return metric \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..eb45643 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1 @@ +# Utils package for SDPose OOD \ No newline at end of file diff --git a/utils/dataset_loader.py b/utils/dataset_loader.py new file mode 100644 index 0000000..f60b3aa --- /dev/null +++ b/utils/dataset_loader.py @@ -0,0 +1,179 @@ +# Author: T. S. Liang @ Rama Alpaca +# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com +# Date: Feb. 2025 +# Description: The dataloader for the training and evaluation protocol. +# License: MIT License (see LICENSE for details) + +from torch.utils.data import DataLoader +from mmengine.dataset import default_collate +from mmengine.registry import init_default_scope +from mmpose.registry import DATASETS + +init_default_scope('mmpose') + +def prepare_dataset(batch_size, dataset_name, dataset_root, mode = "val", num_workers = 4): + + image_size = [768, 1024] ## width x height + sigma = 6 ## sigma is 2 for 256 + scale = 4 + + codec = dict( + type='UDPHeatmap', input_size=(image_size[0], image_size[1]), heatmap_size=(int(image_size[0]/scale), int(image_size[1]/scale)), sigma=sigma) + + train_pipeline = [ + dict(type='LoadImage'), + dict(type='GetBBoxCenterScale'), + dict(type='RandomFlip', direction='horizontal'), + dict(type='RandomHalfBody'), + dict(type='RandomBBoxTransform'), + dict(type='TopdownAffine', input_size=codec['input_size'], use_udp=True), + dict(type='PhotometricDistortion'), + dict(type='CacheRGBTarget'), + dict( + type='Albumentation', + transforms=[ + dict(type='Blur', p=0.1), + dict(type='MedianBlur', p=0.1), + dict( + type='CoarseDropout', + max_holes=1, + max_height=0.4, + max_width=0.4, + min_holes=1, + min_height=0.2, + min_width=0.2, + p=1.0), + ]), + dict(type='GenerateTarget', encoder=codec), + dict(type='PackPoseInputs'), + ] + + val_pipeline = [ + dict(type='LoadImage'), + dict(type='GetBBoxCenterScale'), + dict(type='TopdownAffine', input_size=codec['input_size'], use_udp=True), + dict(type='CacheRGBTarget'), + dict(type='PackPoseInputs'), + ] + + train_cfg = dict( + type='CocoDataset', + data_root=dataset_root, + data_mode='topdown', + ann_file='COCO/annotations/person_keypoints_train2017.json', + data_prefix=dict(img='COCO/train2017/'), + pipeline=train_pipeline, + ) + + if dataset_name == "COCO_OOD": + + print("Running Validation on COCO_OOD.") + + val_cfg = dict( + type='CocoDataset', + data_root=dataset_root, + data_mode='topdown', + ann_file='COCO/annotations/person_keypoints_val2017.json', + data_prefix=dict(img='COCO/val2017oil/'), + test_mode=True, + pipeline=val_pipeline, + bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json', + ) + + elif dataset_name == "Humanart": + + print("Running Validation on Humanart.") + + val_cfg = dict( + type='HumanArtDataset', + data_root=dataset_root, + data_mode='topdown', + ann_file='HumanArt/annotations/validation_humanart.json', + test_mode=True, + pipeline=val_pipeline, + ) + + elif dataset_name == "COCO": + + print("Running Validation on COCO.") + + val_cfg = dict( + type='CocoDataset', + data_root=dataset_root, + data_mode='topdown', + ann_file='COCO/annotations/person_keypoints_val2017.json', + data_prefix=dict(img='COCO/val2017/'), + test_mode=True, + pipeline=val_pipeline, + bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json', + ) + + elif dataset_name == "COCOWholebody": + + print("Running Validation on COCO Wholebody.") + + val_cfg = dict( + type='CocoWholeBodyDataset', + data_root=dataset_root, + data_mode='topdown', + ann_file='COCO/annotations/coco_wholebody_val_v1.0.json', + data_prefix=dict(img='COCO/val2017/'), + test_mode=True, + pipeline=val_pipeline, + bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json', + ) + + elif dataset_name == "COCO-OOD_Wholebody": + + print("Running Validation on COCO-OOD Wholebody.") + + val_cfg = dict( + type='CocoWholeBodyDataset', + data_root=dataset_root, + data_mode='topdown', + ann_file='COCO/annotations/coco_wholebody_val_v1.0.json', + data_prefix=dict(img='COCO/val2017oil/'), + test_mode=True, + pipeline=val_pipeline, + bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json', + ) + + if mode == "train": + + train_ds = DATASETS.build(train_cfg) + val_ds = DATASETS.build(val_cfg) + + train_loader = DataLoader( + train_ds, + batch_size=batch_size, + shuffle=True, + num_workers=num_workers, + collate_fn=default_collate, + pin_memory=True, + ) + + val_loader = DataLoader( + val_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + collate_fn=default_collate, + pin_memory=True, + ) + + return train_loader, val_loader + + elif mode == "val": + + val_ds = DATASETS.build(val_cfg) + + val_loader = DataLoader( + val_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + collate_fn=default_collate, + pin_memory=True, + ) + + return val_loader diff --git a/workflow/ibb_pose.json b/workflow/ibb_pose.json new file mode 100644 index 0000000..ba0da55 --- /dev/null +++ b/workflow/ibb_pose.json @@ -0,0 +1 @@ +{"id":"1bc4458c-a13c-41de-aa48-4bfd5200441b","revision":0,"last_node_id":60,"last_link_id":130,"nodes":[{"id":48,"type":"VHS_LoadVideo","pos":[753.5818705923772,1043.1540168316258],"size":[253.279296875,738.4962614002146],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"video","name":"video","type":"COMBO","widget":{"name":"video"},"link":null},{"localized_name":"force_rate","name":"force_rate","type":"FLOAT","widget":{"name":"force_rate"},"link":null},{"localized_name":"custom_width","name":"custom_width","type":"INT","widget":{"name":"custom_width"},"link":null},{"localized_name":"custom_height","name":"custom_height","type":"INT","widget":{"name":"custom_height"},"link":null},{"localized_name":"frame_load_cap","name":"frame_load_cap","type":"INT","widget":{"name":"frame_load_cap"},"link":null},{"localized_name":"skip_first_frames","name":"skip_first_frames","type":"INT","widget":{"name":"skip_first_frames"},"link":null},{"localized_name":"select_every_nth","name":"select_every_nth","type":"INT","widget":{"name":"select_every_nth"},"link":null},{"localized_name":"format","name":"format","shape":7,"type":"COMBO","widget":{"name":"format"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[]},{"localized_name":"frame_count","name":"frame_count","type":"INT","links":null},{"localized_name":"audio","name":"audio","type":"AUDIO","links":null},{"localized_name":"video_info","name":"video_info","type":"VHS_VIDEOINFO","links":null}],"properties":{"cnr_id":"comfyui-videohelpersuite","ver":"8550981384301e9bc5bfea83e5c2c75258102593","Node name for S&R":"VHS_LoadVideo"},"widgets_values":{"video":"19-工程结算，摇一个-1080P 高清-AVC.mp4","force_rate":0,"custom_width":0,"custom_height":0,"frame_load_cap":10,"skip_first_frames":300,"select_every_nth":1,"format":"AnimateDiff","choose video to upload":"image","videopreview":{"hidden":false,"paused":false,"params":{"filename":"19-工程结算，摇一个-1080P 高清-AVC.mp4","type":"input","format":"video/mp4","force_rate":0,"custom_width":0,"custom_height":0,"frame_load_cap":10,"skip_first_frames":300,"select_every_nth":1}}}},{"id":44,"type":"Note","pos":[1040.1789855910004,1034.998324149427],"size":[210,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["可加载视频或图片"],"color":"#432","bgcolor":"#653"},{"id":59,"type":"Note","pos":[861.5639388065262,777.222864913241],"size":[210,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["可选择使用florence2节点、groundingdino、yolo或均不使用"],"color":"#432","bgcolor":"#653"},{"id":52,"type":"PreviewImage","pos":[1835.548753050627,636.2609251561124],"size":[229.90386295916687,395.228218022673],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":92}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.66","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":50,"type":"Florence2ModelLoader","pos":[837.1175543941788,586.4048382877829],"size":[288.740234375,130],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"lora","name":"lora","shape":7,"type":"PEFTLORA","link":null},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"precision","name":"precision","type":"COMBO","widget":{"name":"precision"},"link":null},{"localized_name":"attention","name":"attention","type":"COMBO","widget":{"name":"attention"},"link":null},{"localized_name":"convert_to_safetensors","name":"convert_to_safetensors","shape":7,"type":"BOOLEAN","widget":{"name":"convert_to_safetensors"},"link":null}],"outputs":[{"localized_name":"florence2_model","name":"florence2_model","type":"FL2MODEL","links":[91]}],"properties":{"cnr_id":"comfyui-florence2","ver":"00b63382966a444a9fefacb65b8deb188d12a458","Node name for S&R":"Florence2ModelLoader","aux_id":"kijai/ComfyUI-Florence2"},"widgets_values":["Florence-2-large-PromptGen-v2.0","fp16","sdpa",false]},{"id":54,"type":"IBBGroundingDinoModelLoader","pos":[1316.0148395732117,1266.5764698982034],"size":[341.905859375,58],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"model_name","name":"model_name","type":"COMBO","widget":{"name":"model_name"},"link":null}],"outputs":[{"localized_name":"GROUNDING_DINO_MODEL","name":"GROUNDING_DINO_MODEL","type":"GROUNDING_DINO_MODEL","links":[]}],"properties":{"cnr_id":"sdpose-ood","ver":"6addae981fb0346a3b9a7cea8d0541a5548817b7","Node name for S&R":"IBBGroundingDinoModelLoader","aux_id":"judian17/ComfyUI-SDPose-OOD"},"widgets_values":["GroundingDINO_SwinT_OGC (694MB)"]},{"id":40,"type":"IBBYOLOModelLoader","pos":[1342.0166034722358,1406.036431383735],"size":[270,58],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"model_name","name":"model_name","type":"COMBO","widget":{"name":"model_name"},"link":null}],"outputs":[{"localized_name":"YOLO_MODEL","name":"YOLO_MODEL","type":"YOLO_MODEL","links":[129]}],"properties":{"cnr_id":"sdpose-ood","ver":"c5953e1e112665340c642387129f1bc50fcf3dfa","Node name for S&R":"IBBYOLOModelLoader","aux_id":"judian17/ComfyUI-SDPose-OOD"},"widgets_values":["real_person_detection.pt"]},{"id":45,"type":"LoadImage","pos":[1035.4005895492703,1151.4013157721954],"size":[270,314],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[126]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.3.66","Node name for S&R":"LoadImage"},"widgets_values":["ComfyUI_00008_.png","image"]},{"id":51,"type":"Florence2Run","pos":[1273.4606572797124,628.5796150419711],"size":[400,369],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":null},{"localized_name":"florence2_model","name":"florence2_model","type":"FL2MODEL","link":91},{"localized_name":"text_input","name":"text_input","type":"STRING","widget":{"name":"text_input"},"link":null},{"localized_name":"task","name":"task","type":"COMBO","widget":{"name":"task"},"link":null},{"localized_name":"fill_mask","name":"fill_mask","type":"BOOLEAN","widget":{"name":"fill_mask"},"link":null},{"localized_name":"keep_model_loaded","name":"keep_model_loaded","shape":7,"type":"BOOLEAN","widget":{"name":"keep_model_loaded"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","shape":7,"type":"INT","widget":{"name":"max_new_tokens"},"link":null},{"localized_name":"num_beams","name":"num_beams","shape":7,"type":"INT","widget":{"name":"num_beams"},"link":null},{"localized_name":"do_sample","name":"do_sample","shape":7,"type":"BOOLEAN","widget":{"name":"do_sample"},"link":null},{"localized_name":"output_mask_select","name":"output_mask_select","shape":7,"type":"STRING","widget":{"name":"output_mask_select"},"link":null},{"localized_name":"seed","name":"seed","shape":7,"type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[92]},{"localized_name":"mask","name":"mask","type":"MASK","links":null},{"localized_name":"caption","name":"caption","type":"STRING","links":[]},{"localized_name":"data","name":"data","type":"JSON","links":[]}],"properties":{"cnr_id":"comfyui-florence2","ver":"00b63382966a444a9fefacb65b8deb188d12a458","Node name for S&R":"Florence2Run","aux_id":"kijai/ComfyUI-Florence2"},"widgets_values":["human","caption_to_phrase_grounding",true,false,1024,3,true,"",2,"fixed",[false,true]]},{"id":32,"type":"VHS_VideoCombine","pos":[2113.4305154601307,1015.6689213997159],"size":[220.5830078125,334],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":null},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":null},{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":null},{"localized_name":"loop_count","name":"loop_count","type":"INT","widget":{"name":"loop_count"},"link":null},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"pingpong","name":"pingpong","type":"BOOLEAN","widget":{"name":"pingpong"},"link":null},{"localized_name":"save_output","name":"save_output","type":"BOOLEAN","widget":{"name":"save_output"},"link":null},{"name":"pix_fmt","type":["yuv420p","yuv420p10le"],"widget":{"name":"pix_fmt"},"link":null},{"name":"crf","type":"INT","widget":{"name":"crf"},"link":null},{"name":"save_metadata","type":"BOOLEAN","widget":{"name":"save_metadata"},"link":null},{"name":"trim_to_audio","type":"BOOLEAN","widget":{"name":"trim_to_audio"},"link":null}],"outputs":[{"localized_name":"Filenames","name":"Filenames","type":"VHS_FILENAMES","links":null}],"properties":{"cnr_id":"comfyui-videohelpersuite","ver":"8550981384301e9bc5bfea83e5c2c75258102593","Node name for S&R":"VHS_VideoCombine","aux_id":"Kosinkadink/ComfyUI-VideoHelperSuite"},"widgets_values":{"frame_rate":30,"loop_count":0,"filename_prefix":"AnimateDiff","format":"video/h264-mp4","pix_fmt":"yuv420p","crf":19,"save_metadata":true,"trim_to_audio":false,"pingpong":false,"save_output":false,"videopreview":{"hidden":false,"paused":false,"params":{"filename":"AnimateDiff_00001.mp4","subfolder":"","type":"temp","format":"video/h264-mp4","frame_rate":30,"workflow":"AnimateDiff_00001.png","fullpath":"C:\\mysoft\\ComfyUI\\temp\\AnimateDiff_00001.mp4"}}}},{"id":39,"type":"IBBPoseModelLoader","pos":[1334.1057814493638,1083.9627438245636],"size":[270,130],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"model_type","name":"model_type","type":"COMBO","widget":{"name":"model_type"},"link":null},{"localized_name":"unet_precision","name":"unet_precision","type":"COMBO","widget":{"name":"unet_precision"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"unload_on_finish","name":"unload_on_finish","type":"BOOLEAN","widget":{"name":"unload_on_finish"},"link":null}],"outputs":[{"localized_name":"IBB_POSE_MODEL","name":"IBB_POSE_MODEL","type":"IBB_POSE_MODEL","links":[122]}],"properties":{"cnr_id":"sdpose-ood","ver":"c5953e1e112665340c642387129f1bc50fcf3dfa","Node name for S&R":"IBBPoseModelLoader","aux_id":"judian17/ComfyUI-SDPose-OOD"},"widgets_values":["WholeBody","bf16","cuda",false]},{"id":46,"type":"PreviewImage","pos":[2408.0297112868234,1094.4467247002833],"size":[140,246],"flags":{},"order":13,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":130}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.66","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":58,"type":"IBBPoseProcessor","pos":[1697.8629829878914,1132.2582805592876],"size":[287.806640625,402],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"ibb_pose_model","name":"ibb_pose_model","type":"IBB_POSE_MODEL","link":122},{"localized_name":"images","name":"images","type":"IMAGE","link":126},{"localized_name":"data_from_florence2","name":"data_from_florence2","shape":7,"type":"JSON","link":null},{"localized_name":"grounding_dino_model","name":"grounding_dino_model","shape":7,"type":"GROUNDING_DINO_MODEL","link":null},{"localized_name":"yolo_model","name":"yolo_model","shape":7,"type":"YOLO_MODEL","link":129},{"localized_name":"score_threshold","name":"score_threshold","type":"FLOAT","widget":{"name":"score_threshold"},"link":null},{"localized_name":"overlay_alpha","name":"overlay_alpha","type":"FLOAT","widget":{"name":"overlay_alpha"},"link":null},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null},{"localized_name":"prompt","name":"prompt","shape":7,"type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"gd_threshold","name":"gd_threshold","shape":7,"type":"FLOAT","widget":{"name":"gd_threshold"},"link":null},{"localized_name":"save_for_editor","name":"save_for_editor","shape":7,"type":"BOOLEAN","widget":{"name":"save_for_editor"},"link":null},{"localized_name":"filename_prefix_edit","name":"filename_prefix_edit","shape":7,"type":"STRING","widget":{"name":"filename_prefix_edit"},"link":null},{"localized_name":"keep_face","name":"keep_face","shape":7,"type":"BOOLEAN","widget":{"name":"keep_face"},"link":null},{"localized_name":"keep_hands","name":"keep_hands","shape":7,"type":"BOOLEAN","widget":{"name":"keep_hands"},"link":null},{"localized_name":"keep_feet","name":"keep_feet","shape":7,"type":"BOOLEAN","widget":{"name":"keep_feet"},"link":null},{"localized_name":"scale_for_xinsr","name":"scale_for_xinsr","shape":7,"type":"BOOLEAN","widget":{"name":"scale_for_xinsr"},"link":null},{"localized_name":"pose_scale_factor","name":"pose_scale_factor","shape":7,"type":"FLOAT","widget":{"name":"pose_scale_factor"},"link":null}],"outputs":[{"localized_name":"images","name":"images","type":"IMAGE","links":[130]},{"localized_name":"pose_keypoint","name":"pose_keypoint","type":"POSE_KEYPOINT","links":null}],"properties":{"cnr_id":"sdpose-ood","ver":"6addae981fb0346a3b9a7cea8d0541a5548817b7","Node name for S&R":"IBBPoseProcessor","aux_id":"judian17/ComfyUI-SDPose-OOD"},"widgets_values":[0.3,0.6,1,"person .",0.3,false,"poses/ibb_pose",true,true,true,true,1]},{"id":18,"type":"Note","pos":[1295.668791203465,1557.838806386414],"size":[776.3567180428074,291.2611697064141],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["body 模型为17个点的仅身体姿态；wholebody模型为包括面部和手指的133点的姿态\n30系以上显卡推荐bf16，更早的使用fp16。16精度速度不变显存占用降低。\ncpu推理仅可用fp32\n\nscore threshold控制识别效果\noverlay alpha控制结果图背景可见度。1代表纯黑背景的仅姿态图\nbatch_size最初设计目的为并行运行增加速度，后来发现并行了也没什么区别。可能类似4090这种高算力的能有区别吧。但是修改太麻烦于是保留。正常用1即可。\n\ndata_from_florence2接受来自https://github.com/kijai/ComfyUI-Florence2.git\n节点的“data输出”，需注意需切换识别模式为 caption_to_phrase_grounding,并使用英文提示词指定想要识别的对象。可用people、human等识别所有、man、woman等仅识别特定性别，或其他特征的人物。\n\ngrounding_dion_model链接groundingdino model节点，以使用groundingdino来识别图中人物。在节点的prompt输入里输入groundingdino所识别的对象。参数gd_threshold控制groundingdino识别效果。速度最慢\n\nyolo_model接受load yolo model节点输入，使用yolo来生成bbox。速度最快\n\n\n以上三种方式均选其中一种，都不使用则直接识别全图，多人时可能识别不全。\n\n\nsave for editor选项开启后，可在output/poses下找到pose_edit前缀的姿态json文件，用于https://github.com/judian17/ComfyUI-OpenPose-Editor-jd 此节点导入并编辑姿态。\n注意保存的json格式与输出的字符串类型的pose_json不太一样，后者为项目自带格式，无法用于上述节点的导入编辑\n\n当使用wholebody模型时，可通过keep face、keep hands、keep feet参数控制是否输出面部、手指、脚的点。使用body模型时无用。\n\nscale_for_xinsr：sdxl的union模型使用姿态控制时需要使用打开此模式获取肢体膨胀后的姿态图。\n\npose_scale_factor:当原图接近512尺寸时此值为1，姿态图比例较正常。当原图为1024左右时此值设置为2比例正常，否则肢体会偏细，关节偏小。"],"color":"#432","bgcolor":"#653"}],"links":[[91,50,0,51,1,"FL2MODEL"],[92,51,0,52,0,"IMAGE"],[122,39,0,58,0,"IBB_POSE_MODEL"],[126,45,0,58,1,"IMAGE"],[129,40,0,58,4,"YOLO_MODEL"],[130,58,0,46,0,"IMAGE"]],"groups":[],"config":{},"extra":{"ds":{"scale":1.2284597357367297,"offset":[-1072.875069139725,-1329.7873272689385]},"workflowRendererVersion":"LG"},"version":0.4} \ No newline at end of file