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>
This commit is contained in:
copilot-swe-agent[bot]
co-authored by
kukulu00
kukulu00
parent
bb6f886d33
commit
559f45a67b
@@ -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
|
||||
|
||||
|
||||
+269
-36
@@ -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
|
||||
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"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user