From 559f45a67b4f8ffadcf0ec1cf72aceb51f3fe360 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 22 Apr 2026 08:26:45 +0000 Subject: [PATCH] Add fallback SDPose heatmap head Agent-Logs-Url: https://github.com/IBB666/IBB_POSE/sessions/be2c1a27-aacc-43a8-adc6-0bb4dbae296e Co-authored-by: kukulu00 <1170651+kukulu00@users.noreply.github.com> --- README.md | 1 + models/HeatmapHead.py | 305 +++++++++++++++++++++++++++++++++++++----- 2 files changed, 270 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 2627a8b..c24807c 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,7 @@ Inspired by [SDPose-OOD](https://github.com/judian17/ComfyUI-SDPose-OOD). python -m pip cache purge ``` - The node now keeps heavy pose backends optional and only imports them when the selected mode needs them. +- The ComfyUI runtime no longer needs the OpenMMLab stack just to load the SDPose decoder. `mmpose` / `mmcv` / `mmengine` are only needed for the standalone evaluation scripts under `/eval`. ### Manual model placement diff --git a/models/HeatmapHead.py b/models/HeatmapHead.py index b5268b7..57bc725 100644 --- a/models/HeatmapHead.py +++ b/models/HeatmapHead.py @@ -7,58 +7,291 @@ # 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 +from __future__ import annotations -def get_heatmap_head(mode = "body"): +import warnings +from types import SimpleNamespace + +import numpy as np +import torch +from torch import nn + +try: + from mmpose.registry import MODELS +except Exception as _mmpose_import_error: + MODELS = None +else: + _mmpose_import_error = None + + +def _build_head_cfg(mode="body"): if mode == "body": - image_size = [768, 1024] ## width x height - sigma = 6 ## sigma is 2 for 256 + 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 + 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 + else: + raise ValueError(f"Unsupported heatmap head mode: {mode}") - 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 + 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, + ) + + return dict( + type="HeatmapHead", + in_channels=embed_dim, + out_channels=num_keypoints, + deconv_out_channels=(embed_dim,), + deconv_kernel_sizes=(4,), + conv_out_channels=(embed_dim,), + conv_kernel_sizes=(1,), + loss=dict(type="KeypointMSELoss", use_target_weight=True), + decoder=codec, + ) + + +class _CompatInstanceData(SimpleNamespace): + def set_field(self, value, name): + setattr(self, name, value) + + def all_items(self): + return self.__dict__.items() + + +class _FallbackUDPHeatmap: + def __init__(self, input_size, heatmap_size, sigma=2.0): + self.input_size = np.asarray(input_size, dtype=np.float32) + self.heatmap_size = np.asarray(heatmap_size, dtype=np.float32) + self.sigma = sigma + + def decode(self, encoded): + heatmaps = encoded.astype(np.float32, copy=False) + num_keypoints, height, width = heatmaps.shape + + flattened = heatmaps.reshape(num_keypoints, -1) + flat_indices = flattened.argmax(axis=1) + scores = flattened[np.arange(num_keypoints), flat_indices] + + xs = (flat_indices % width).astype(np.float32) + ys = (flat_indices // width).astype(np.float32) + + for idx, (x, y) in enumerate(zip(xs.astype(np.int32), ys.astype(np.int32))): + if 0 < x < width - 1: + xs[idx] += np.sign(heatmaps[idx, y, x + 1] - heatmaps[idx, y, x - 1]) * 0.25 + if 0 < y < height - 1: + ys[idx] += np.sign(heatmaps[idx, y + 1, x] - heatmaps[idx, y - 1, x]) * 0.25 + + scale = self.input_size / np.maximum(self.heatmap_size - 1.0, 1.0) + keypoints = np.stack((xs, ys), axis=-1) * scale + + return keypoints[None], scores[None].astype(np.float32) + + +class _FallbackHeatmapHead(nn.Module): + _version = 2 + + def __init__( + self, + in_channels, + out_channels, + deconv_out_channels=(256, 256, 256), + deconv_kernel_sizes=(4, 4, 4), + conv_out_channels=None, + conv_kernel_sizes=None, + final_layer=None, + loss=None, + decoder=None, + init_cfg=None, + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.loss_module = None + self.decoder = _FallbackUDPHeatmap(**decoder) if decoder is not None else None + + if deconv_out_channels: + self.deconv_layers = self._make_deconv_layers( + in_channels=in_channels, + layer_out_channels=deconv_out_channels, + layer_kernel_sizes=deconv_kernel_sizes, + ) + in_channels = deconv_out_channels[-1] + else: + self.deconv_layers = nn.Identity() + + if conv_out_channels: + self.conv_layers = self._make_conv_layers( + in_channels=in_channels, + layer_out_channels=conv_out_channels, + layer_kernel_sizes=conv_kernel_sizes, + ) + in_channels = conv_out_channels[-1] + else: + self.conv_layers = nn.Identity() + + final_layer = final_layer or {"kernel_size": 1} + self.final_layer = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=final_layer.get("kernel_size", 1), + stride=final_layer.get("stride", 1), + padding=final_layer.get("padding", 0), ) - head = MODELS.build(head_cfg) - - return head \ No newline at end of file + self._register_load_state_dict_pre_hook(self._load_state_dict_pre_hook) + + @staticmethod + def _make_conv_layers(in_channels, layer_out_channels, layer_kernel_sizes): + layers = [] + for out_channels, kernel_size in zip(layer_out_channels, layer_kernel_sizes): + padding = (kernel_size - 1) // 2 + layers.append( + nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=1, + padding=padding, + ) + ) + layers.append(nn.BatchNorm2d(num_features=out_channels)) + layers.append(nn.ReLU(inplace=True)) + in_channels = out_channels + return nn.Sequential(*layers) + + @staticmethod + def _make_deconv_layers(in_channels, layer_out_channels, layer_kernel_sizes): + layers = [] + for out_channels, kernel_size in zip(layer_out_channels, layer_kernel_sizes): + if kernel_size == 4: + padding = 1 + output_padding = 0 + elif kernel_size == 3: + padding = 1 + output_padding = 1 + elif kernel_size == 2: + padding = 0 + output_padding = 0 + else: + raise ValueError(f"Unsupported kernel size {kernel_size} for deconv layers") + + layers.append( + nn.ConvTranspose2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=2, + padding=padding, + output_padding=output_padding, + bias=False, + ) + ) + layers.append(nn.BatchNorm2d(num_features=out_channels)) + layers.append(nn.ReLU(inplace=True)) + in_channels = out_channels + return nn.Sequential(*layers) + + def forward(self, feats): + x = feats[-1] if isinstance(feats, (list, tuple)) else feats + x = self.deconv_layers(x) + x = self.conv_layers(x) + x = self.final_layer(x) + return x + + def decode(self, batch_heatmaps): + heatmaps_np = batch_heatmaps.detach().cpu().float().numpy() + preds = [] + for heatmaps in heatmaps_np: + keypoints, scores = self.decoder.decode(heatmaps) + preds.append(_CompatInstanceData(keypoints=keypoints, keypoint_scores=scores)) + return preds + + def predict(self, feats, batch_data_samples=None, test_cfg=None): + test_cfg = test_cfg or {} + + if test_cfg.get("flip_test", False): + raise RuntimeError( + "The lightweight fallback heatmap head does not support MMPose flip_test. " + "Install mmpose/mmcv for evaluation workflows." + ) + + batch_heatmaps = self.forward(feats) + preds = self.decode(batch_heatmaps) + + if test_cfg.get("output_heatmaps", False): + pred_fields = [SimpleNamespace(heatmaps=hm) for hm in batch_heatmaps.detach()] + return preds, pred_fields + + return preds + + def _load_state_dict_pre_hook(self, state_dict, prefix, local_meta, *args, **kwargs): + version = local_meta.get("version", None) + if version and version >= self._version: + return + + keys = list(state_dict.keys()) + for key in keys: + if not key.startswith(prefix): + continue + + value = state_dict.pop(key) + short_key = key[len(prefix):] + key_parts = short_key.split(".") + + if key_parts[0] == "final_layer": + if len(key_parts) == 3: + idx = int(key_parts[1]) + if idx < len(self.conv_layers): + new_key = "conv_layers." + ".".join(key_parts[1:]) + else: + new_key = "final_layer." + key_parts[2] + else: + new_key = short_key + else: + new_key = short_key + + state_dict[prefix + new_key] = value + + +def get_heatmap_head(mode="body"): + head_cfg = _build_head_cfg(mode=mode) + + if MODELS is not None: + return MODELS.build(head_cfg) + + warnings.warn( + "IBB_POSE: mmpose/mmcv is not available " + f"({_mmpose_import_error}). Falling back to the built-in heatmap head " + "for ComfyUI inference. Install mmpose/mmcv/mmengine only if you need the evaluation scripts.", + RuntimeWarning, + stacklevel=2, + ) + + return _FallbackHeatmapHead( + in_channels=head_cfg["in_channels"], + out_channels=head_cfg["out_channels"], + deconv_out_channels=head_cfg["deconv_out_channels"], + deconv_kernel_sizes=head_cfg["deconv_kernel_sizes"], + conv_out_channels=head_cfg["conv_out_channels"], + conv_kernel_sizes=head_cfg["conv_kernel_sizes"], + loss=head_cfg["loss"], + decoder=head_cfg["decoder"], + )