489 lines
18 KiB
Python
489 lines
18 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
|
|
import pytest
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from src.Comfyui_Fd_Nodes import zhiyi_dwpose_detect_node as dwpose_module
|
|
from src.Comfyui_Fd_Nodes.utils import image_payload as payload_module
|
|
from src.Comfyui_Fd_Nodes.v2.controlnet_aux_v2_node import ZhiYiDWPoseDetectNodeV2
|
|
from src.Comfyui_Fd_Nodes.config import CUSTOM_SERVICE_URL_PRESET
|
|
from src.Comfyui_Fd_Nodes.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
|
from src.Comfyui_Fd_Nodes.zhiyi_dwpose_detect_node import (
|
|
ZhiYiDWPoseDetectNode,
|
|
_health_url,
|
|
_normalize_pose_url,
|
|
)
|
|
|
|
|
|
def _image_data_url(size=(2, 2), color=(255, 0, 0)):
|
|
image = Image.new("RGB", size, color)
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
return f"data:image/png;base64,{encoded}"
|
|
|
|
|
|
class DummyResponse:
|
|
def __init__(self, status_code=200, data=None, text=""):
|
|
self.status_code = status_code
|
|
self._data = data
|
|
self.text = text
|
|
|
|
def json(self):
|
|
if isinstance(self._data, Exception):
|
|
raise self._data
|
|
return self._data
|
|
|
|
|
|
def _pose_response(request_id="pose-1", size=(2, 2), people_count=1):
|
|
return {
|
|
"image_size": {"width": size[0], "height": size[1]},
|
|
"inference_size": {"width": size[0], "height": size[1]},
|
|
"pose_image": _image_data_url(size=size, color=(0, 255, 0)),
|
|
"openpose_json": {
|
|
"people": [{"pose_keypoints_2d": [1, 2, 0.9]}],
|
|
"canvas_width": size[0],
|
|
"canvas_height": size[1],
|
|
},
|
|
"people_count": people_count,
|
|
"request_id": request_id,
|
|
"model_name": "DWPose",
|
|
"detector_filename": "yolox_l.onnx",
|
|
"pose_filename": "dw-ll_ucoco_384.onnx",
|
|
"device": "cpu",
|
|
"inference_ms": 12.3,
|
|
"total_ms": 45.6,
|
|
}
|
|
|
|
|
|
def test_dwpose_node_metadata():
|
|
input_types = ZhiYiDWPoseDetectNode.INPUT_TYPES()
|
|
|
|
assert input_types["required"]["image"][0] == "IMAGE"
|
|
assert input_types["required"]["service_url"][1]["default"].endswith("/v1/pose")
|
|
assert input_types["required"]["detect_face"][1]["default"] is False
|
|
assert "health_check" in input_types["optional"]
|
|
assert "node_switch" in input_types["optional"]
|
|
assert list(input_types["optional"])[-2:] == ["service_url_preset", "image_delivery"]
|
|
preset_options, preset_config = input_types["optional"]["service_url_preset"]
|
|
assert preset_options == [CUSTOM_SERVICE_URL_PRESET, "K8s 灰度", "海外接入", "10.1.0.230"]
|
|
assert preset_config["default"] == CUSTOM_SERVICE_URL_PRESET
|
|
assert ZhiYiDWPoseDetectNode.RETURN_TYPES == ("IMAGE", "POSE_KEYPOINT", "JSON")
|
|
assert ZhiYiDWPoseDetectNode.RETURN_NAMES == ("POSE_IMAGE", "POSE_KEYPOINT", "INFO")
|
|
assert ZhiYiDWPoseDetectNode.CATEGORY == "知衣/姿态检测"
|
|
assert NODE_CLASS_MAPPINGS["ZhiYiDWPoseDetectNode"] is ZhiYiDWPoseDetectNode
|
|
assert NODE_DISPLAY_NAME_MAPPINGS["ZhiYiDWPoseDetectNode"] == "知衣-DWPose姿态检测"
|
|
|
|
|
|
def test_normalize_pose_url_variants():
|
|
assert (
|
|
_normalize_pose_url("model-api-dwpose-svc.online-server-gray:8001")
|
|
== "http://model-api-dwpose-svc.online-server-gray:8001/v1/pose"
|
|
)
|
|
assert _normalize_pose_url("http://example.com:8001") == "http://example.com:8001/v1/pose"
|
|
assert _normalize_pose_url("https://example.com/v1/pose/") == "https://example.com/v1/pose"
|
|
assert _normalize_pose_url("https://example.com/api") == "https://example.com/api/v1/pose"
|
|
assert _health_url("https://example.com/v1/pose") == "https://example.com/health"
|
|
assert _health_url("https://example.com/api/v1/pose") == "https://example.com/api/health"
|
|
|
|
|
|
def test_tensor_data_url_roundtrip():
|
|
node = ZhiYiDWPoseDetectNode()
|
|
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
image[:, :, :, 0] = 1.0
|
|
|
|
data_url = node._tensor_to_png_data_url(image)
|
|
decoded = node._decode_png_data_url(data_url, "pose_image")
|
|
tensor = node._pil_to_image_tensor(decoded)
|
|
|
|
assert data_url.startswith("data:image/png;base64,")
|
|
assert decoded.size == (2, 2)
|
|
assert tensor.shape == (1, 2, 2, 3)
|
|
assert torch.equal(tensor[0, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
|
|
|
|
|
|
def test_detect_pose_posts_payload_and_decodes_outputs(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
images = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
calls = []
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
calls.append((method, url, timeout, kwargs))
|
|
assert method == "POST"
|
|
assert url == "http://example.com:8001/v1/pose"
|
|
assert kwargs["headers"] == {"Content-Type": "application/json"}
|
|
payload = kwargs["json"]
|
|
assert payload["image"].startswith("data:image/png;base64,")
|
|
assert payload["resolution"] == 768
|
|
assert payload["detect_body"] is True
|
|
assert payload["detect_hand"] is False
|
|
assert payload["detect_face"] is False
|
|
assert payload["xinsr_stick_scaling"] is True
|
|
assert payload["upscale_method"] == "INTER_AREA"
|
|
assert payload["return_pose_image"] is True
|
|
assert payload["return_openpose_json"] is True
|
|
assert payload["image_format"] == "png_base64"
|
|
return DummyResponse(data=_pose_response(request_id="pose-1", size=(2, 2)))
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
pose_image, pose_keypoint, info_text = node.detect_pose(
|
|
image=images,
|
|
service_url="example.com:8001",
|
|
resolution=768,
|
|
detect_body=True,
|
|
detect_hand=False,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=True,
|
|
upscale_method="INTER_AREA",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
)
|
|
|
|
info = json.loads(info_text)
|
|
assert len(calls) == 1
|
|
assert pose_image.shape == (1, 2, 2, 3)
|
|
assert torch.equal(pose_image[0, 0, 0], torch.tensor([0.0, 1.0, 0.0]))
|
|
assert pose_keypoint[0]["people"][0]["pose_keypoints_2d"] == [1, 2, 0.9]
|
|
assert info["service_url"] == "http://example.com:8001/v1/pose"
|
|
assert info["resolution"] == 768
|
|
assert info["detect_hand"] is False
|
|
assert info["items"][0]["request_id"] == "pose-1"
|
|
assert info["items"][0]["people_count"] == 1
|
|
|
|
|
|
def test_detect_pose_preset_overrides_custom_url(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
assert method == "POST"
|
|
assert url == "http://10.1.0.230:8003/v1/pose"
|
|
return DummyResponse(data=_pose_response())
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
_, _, info_text = node.detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://custom.example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
service_url_preset="10.1.0.230",
|
|
)
|
|
|
|
assert json.loads(info_text)["service_url"] == "http://10.1.0.230:8003/v1/pose"
|
|
|
|
|
|
def test_detect_pose_overseas_preset_resolves_gateway_url(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
assert url == (
|
|
"http://fd-oversease-access-hz.external.sometimes.cn"
|
|
"/gpu/online-server-gray/model-api-dwpose/v1/pose"
|
|
)
|
|
return DummyResponse(data=_pose_response())
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
_, _, info_text = node.detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://custom.example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
service_url_preset="海外接入",
|
|
)
|
|
|
|
assert json.loads(info_text)["service_url"] == (
|
|
"http://fd-oversease-access-hz.external.sometimes.cn"
|
|
"/gpu/online-server-gray/model-api-dwpose/v1/pose"
|
|
)
|
|
|
|
|
|
def test_detect_pose_rejects_unknown_preset_but_node_switch_still_short_circuits(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
|
|
with pytest.raises(RuntimeError, match="未知的 service_url_preset"):
|
|
node.detect_pose(
|
|
image=image,
|
|
service_url="https://custom.example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
service_url_preset="不存在",
|
|
)
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", lambda *args, **kwargs: pytest.fail("不应发出请求"))
|
|
result_image, keypoints, info_text = node.detect_pose(
|
|
image=image,
|
|
service_url="https://custom.example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
node_switch=1,
|
|
service_url_preset="不存在",
|
|
)
|
|
assert torch.equal(result_image, image)
|
|
assert keypoints == []
|
|
assert json.loads(info_text)["skipped"] is True
|
|
|
|
|
|
def test_detect_pose_batch_keeps_order_and_resizes_pose_image(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
images = torch.zeros((2, 2, 3, 3), dtype=torch.float32)
|
|
call_count = 0
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
return DummyResponse(data=_pose_response(request_id=f"pose-{call_count}", size=(1, 1), people_count=call_count))
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
pose_image, pose_keypoint, info_text = node.detect_pose(
|
|
image=images,
|
|
service_url="https://example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
)
|
|
|
|
info = json.loads(info_text)
|
|
assert call_count == 2
|
|
assert pose_image.shape == (2, 2, 3, 3)
|
|
assert len(pose_keypoint) == 2
|
|
assert info["items"][0]["request_id"] == "pose-1"
|
|
assert info["items"][1]["request_id"] == "pose-2"
|
|
|
|
|
|
def test_node_switch_returns_original_image_without_request(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
image = torch.ones((1, 2, 2, 3), dtype=torch.float32)
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
raise AssertionError("request should not be called")
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
pose_image, pose_keypoint, info_text = node.detect_pose(
|
|
image=image,
|
|
service_url="https://example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=True,
|
|
node_switch=1,
|
|
)
|
|
|
|
assert torch.equal(pose_image, image)
|
|
assert pose_keypoint == []
|
|
assert json.loads(info_text)["skipped"] is True
|
|
|
|
|
|
def test_health_check_uses_service_root(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
calls = []
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
calls.append((method, url))
|
|
if method == "GET":
|
|
return DummyResponse(data={"status": "ok", "loaded": True})
|
|
return DummyResponse(data=_pose_response(request_id="pose-health"))
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
node.detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=True,
|
|
node_switch=0,
|
|
)
|
|
|
|
assert calls[0] == ("GET", "https://example.com/health")
|
|
assert calls[1] == ("POST", "https://example.com/v1/pose")
|
|
|
|
|
|
def test_health_check_rejects_unloaded(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
return DummyResponse(data={"status": "loading_failed", "loaded": False})
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
with pytest.raises(RuntimeError, match="服务未就绪"):
|
|
node.detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=True,
|
|
node_switch=0,
|
|
)
|
|
|
|
|
|
def test_post_pose_uses_service_error_message():
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
return DummyResponse(
|
|
status_code=500,
|
|
data={"error": {"code": "MODEL_UNAVAILABLE", "message": "not loaded", "request_id": "req-err"}},
|
|
)
|
|
|
|
original_request = dwpose_module.requests.request
|
|
dwpose_module.requests.request = fake_request
|
|
try:
|
|
with pytest.raises(RuntimeError, match="MODEL_UNAVAILABLE: not loaded: request_id=req-err"):
|
|
node._post_pose("https://example.com/v1/pose", {"image": "x"}, timeout=30)
|
|
finally:
|
|
dwpose_module.requests.request = original_request
|
|
|
|
|
|
def test_post_pose_rejects_non_json_response(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
return DummyResponse(data=ValueError("not json"), text="<html>bad</html>")
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
with pytest.raises(RuntimeError, match="响应不是 JSON"):
|
|
node._post_pose("https://example.com/v1/pose", {"image": "x"}, timeout=30)
|
|
|
|
|
|
def test_post_pose_requires_pose_image(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
data = _pose_response()
|
|
data["pose_image"] = None
|
|
return DummyResponse(data=data)
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
with pytest.raises(RuntimeError, match="缺少 pose_image"):
|
|
node._post_pose("https://example.com/v1/pose", {"image": "x"}, timeout=30)
|
|
|
|
|
|
def test_timeout_is_normalized(monkeypatch):
|
|
node = ZhiYiDWPoseDetectNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
raise requests.exceptions.Timeout("read timed out")
|
|
|
|
monkeypatch.setattr(dwpose_module.requests, "request", fake_request)
|
|
|
|
with pytest.raises(RuntimeError, match="TIMEOUT"):
|
|
node.detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://example.com/v1/pose",
|
|
resolution=512,
|
|
detect_body=True,
|
|
detect_hand=True,
|
|
detect_face=False,
|
|
xinsr_stick_scaling=False,
|
|
upscale_method="INTER_CUBIC",
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("node_class", [ZhiYiDWPoseDetectNode, ZhiYiDWPoseDetectNodeV2])
|
|
def test_url_delivery_uses_oss_and_keeps_response_format(monkeypatch, node_class):
|
|
calls, uploads = [], []
|
|
def upload(path, data):
|
|
uploads.append(data)
|
|
return "https://bucket.oss-cn-hangzhou.aliyuncs.com/pose.png"
|
|
def request(method, url, timeout, **kwargs):
|
|
assert kwargs["headers"]["Authorization"] == "Bearer pose-key"
|
|
if method == "GET":
|
|
return DummyResponse(data={"status": "ok", "loaded": True})
|
|
calls.append(kwargs["json"])
|
|
return DummyResponse(data=_pose_response())
|
|
monkeypatch.setattr(payload_module, "upload_bytes_to_oss", upload)
|
|
monkeypatch.setattr(dwpose_module, "FD_DWPOSE_API_KEY", "pose-key")
|
|
monkeypatch.setattr(dwpose_module.requests, "request", request)
|
|
extra = {"env": "custom", "return_pose_image": True, "return_openpose_json": True} if node_class is ZhiYiDWPoseDetectNodeV2 else {}
|
|
result = node_class().detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3)), service_url="https://example.com/v1/pose",
|
|
resolution=512, detect_body=True, detect_hand=True, detect_face=False,
|
|
xinsr_stick_scaling=False, upscale_method="INTER_CUBIC", max_concurrency=1,
|
|
timeout=120, health_check=True, image_delivery="url", **extra,
|
|
)
|
|
assert len(uploads) == len(calls) == 1
|
|
assert calls[0]["image"] == "https://bucket.oss-cn-hangzhou-internal.aliyuncs.com/pose.png"
|
|
assert calls[0]["image_format"] == "png_base64"
|
|
assert result[0].shape == (1, 2, 2, 3)
|
|
assert json.loads(result[2])["image_delivery"] == "url"
|
|
|
|
|
|
def test_url_delivery_node_switch_does_not_upload(monkeypatch):
|
|
def unexpected(*args, **kwargs):
|
|
pytest.fail("跳过节点不应上传或调用服务")
|
|
monkeypatch.setattr(payload_module, "upload_bytes_to_oss", unexpected)
|
|
monkeypatch.setattr(dwpose_module.requests, "request", unexpected)
|
|
result = ZhiYiDWPoseDetectNode().detect_pose(
|
|
image=torch.zeros((1, 2, 2, 3)), service_url="https://example.com/v1/pose",
|
|
resolution=512, detect_body=True, detect_hand=True, detect_face=False,
|
|
xinsr_stick_scaling=False, upscale_method="INTER_CUBIC", max_concurrency=1,
|
|
timeout=120, node_switch=1, image_delivery="url",
|
|
)
|
|
assert json.loads(result[2])["skipped"] is True
|