497 lines
19 KiB
Python
497 lines
19 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_rmbg_segment_node as rmbg_module
|
|
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_rmbg_segment_node import (
|
|
BODY_CLASSES,
|
|
CLOTHES_CLASSES,
|
|
FASHION_CLASSES,
|
|
ZhiYiBodySegmentNode,
|
|
ZhiYiClothesSegmentNode,
|
|
ZhiYiFashionSegmentNode,
|
|
ZhiYiRMBGNode,
|
|
_health_url,
|
|
_normalize_rmbg_segment_url,
|
|
)
|
|
|
|
|
|
def _image_data_url(size=(2, 2), color=(0, 255, 0), mode="RGB"):
|
|
image = Image.new(mode, size, color)
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
return f"data:image/png;base64,{encoded}"
|
|
|
|
|
|
def _mask_data_url(size=(2, 2), white_pixels=None):
|
|
image = Image.new("L", size, 0)
|
|
for xy in white_pixels or []:
|
|
image.putpixel(xy, 255)
|
|
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
|
|
|
|
|
|
_MISSING = object()
|
|
|
|
|
|
def _segment_response(
|
|
*,
|
|
request_id="seg-1",
|
|
size=(2, 2),
|
|
processed_image=_MISSING,
|
|
mask=None,
|
|
mask_image=None,
|
|
classes=None,
|
|
preprocessor="ClothesSegmentPreprocessor",
|
|
):
|
|
return {
|
|
"image_size": {"width": size[0], "height": size[1]},
|
|
"inference_size": {"width": size[0], "height": size[1]},
|
|
"output_size": {"width": size[0], "height": size[1]},
|
|
"processed_image": _image_data_url(size=size, color=(0, 255, 0)) if processed_image is _MISSING else processed_image,
|
|
"mask": mask if mask is not None else _mask_data_url(size=size, white_pixels=[(0, 0), (1, 1)]),
|
|
"mask_image": mask_image,
|
|
"classes": classes or [],
|
|
"request_id": request_id,
|
|
"preprocessor": preprocessor,
|
|
"model_name": "SegFormer Clothes",
|
|
"model_type": "semantic_segmentation",
|
|
"model_repo": "1038lab/segformer_clothes",
|
|
"model_filename": "model.safetensors",
|
|
"model_path": "/models/dwpose/1038lab/segformer_clothes/model.safetensors",
|
|
"device": "cpu",
|
|
"inference_ms": 12.3,
|
|
"total_ms": 45.6,
|
|
}
|
|
|
|
|
|
def test_rmbg_and_segment_node_metadata():
|
|
rmbg_inputs = ZhiYiRMBGNode.INPUT_TYPES()
|
|
clothes_inputs = ZhiYiClothesSegmentNode.INPUT_TYPES()
|
|
fashion_inputs = ZhiYiFashionSegmentNode.INPUT_TYPES()
|
|
body_inputs = ZhiYiBodySegmentNode.INPUT_TYPES()
|
|
|
|
assert rmbg_inputs["required"]["service_url"][1]["default"].endswith("/v1/rmbg")
|
|
assert rmbg_inputs["required"]["process_res"][1]["default"] == 1024
|
|
assert rmbg_inputs["required"]["sensitivity"][1]["default"] == 1.0
|
|
assert rmbg_inputs["required"]["refine_foreground"][1]["default"] is False
|
|
assert rmbg_inputs["required"]["return_mask_image"][1]["default"] is False
|
|
assert clothes_inputs["required"]["service_url"][1]["default"].endswith("/v1/segment/clothes")
|
|
assert fashion_inputs["required"]["service_url"][1]["default"].endswith("/v1/segment/fashion")
|
|
assert body_inputs["required"]["service_url"][1]["default"].endswith("/v1/segment/body")
|
|
assert list(body_inputs["required"]) == [
|
|
"image",
|
|
"service_url",
|
|
"classes",
|
|
"process_res",
|
|
"mask_blur",
|
|
"mask_offset",
|
|
"invert_output",
|
|
"background",
|
|
"background_color",
|
|
"return_image",
|
|
"return_mask",
|
|
"return_mask_image",
|
|
"max_concurrency",
|
|
"timeout",
|
|
]
|
|
assert list(body_inputs["optional"]) == ["health_check", "node_switch", "service_url_preset"]
|
|
assert body_inputs["required"]["service_url"][0] == "STRING"
|
|
assert body_inputs["required"]["classes"][0] == "STRING"
|
|
assert "忽略 process_res" in body_inputs["required"]["process_res"][1]["tooltip"]
|
|
|
|
assert len(CLOTHES_CLASSES) == 18
|
|
assert len(FASHION_CLASSES) == 47
|
|
assert len(BODY_CLASSES) == 12
|
|
assert "shirt, blouse" in FASHION_CLASSES
|
|
assert "top, t-shirt, sweatshirt" in FASHION_CLASSES
|
|
assert "Torso-skin" in BODY_CLASSES
|
|
|
|
for node_cls in (ZhiYiRMBGNode, ZhiYiClothesSegmentNode, ZhiYiFashionSegmentNode, ZhiYiBodySegmentNode):
|
|
assert node_cls.RETURN_TYPES == ("IMAGE", "MASK", "IMAGE", "JSON")
|
|
assert node_cls.RETURN_NAMES == ("IMAGE", "MASK", "MASK_IMAGE", "INFO")
|
|
input_types = node_cls.INPUT_TYPES()
|
|
assert list(input_types["optional"])[-1] == "service_url_preset"
|
|
preset_options, preset_config = input_types["optional"]["service_url_preset"]
|
|
assert preset_options == [CUSTOM_SERVICE_URL_PRESET, "10.1.0.230"]
|
|
assert preset_config["default"] == CUSTOM_SERVICE_URL_PRESET
|
|
|
|
assert ZhiYiRMBGNode.CATEGORY == "知衣/抠图"
|
|
assert ZhiYiClothesSegmentNode.CATEGORY == "知衣/语义分割"
|
|
assert NODE_CLASS_MAPPINGS["ZhiYiRMBGNode"] is ZhiYiRMBGNode
|
|
assert NODE_CLASS_MAPPINGS["ZhiYiClothesSegmentNode"] is ZhiYiClothesSegmentNode
|
|
assert NODE_CLASS_MAPPINGS["ZhiYiFashionSegmentNode"] is ZhiYiFashionSegmentNode
|
|
assert NODE_CLASS_MAPPINGS["ZhiYiBodySegmentNode"] is ZhiYiBodySegmentNode
|
|
assert NODE_DISPLAY_NAME_MAPPINGS["ZhiYiRMBGNode"] == "知衣-RMBG2.0背景去除"
|
|
assert NODE_DISPLAY_NAME_MAPPINGS["ZhiYiClothesSegmentNode"] == "知衣-衣物语义分割"
|
|
assert NODE_DISPLAY_NAME_MAPPINGS["ZhiYiFashionSegmentNode"] == "知衣-时尚单品分割"
|
|
assert NODE_DISPLAY_NAME_MAPPINGS["ZhiYiBodySegmentNode"] == "知衣-身体部位分割"
|
|
|
|
|
|
def test_normalize_rmbg_segment_url_variants():
|
|
assert _normalize_rmbg_segment_url("10.1.0.230:8003", "http://unused/v1/rmbg", "rmbg") == "http://10.1.0.230:8003/v1/rmbg"
|
|
assert _normalize_rmbg_segment_url("https://example.com/api", "http://unused/v1/rmbg", "rmbg") == "https://example.com/api/v1/rmbg"
|
|
assert (
|
|
_normalize_rmbg_segment_url("https://example.com/api/v1", "http://unused/v1/segment/clothes", "segment/clothes")
|
|
== "https://example.com/api/v1/segment/clothes"
|
|
)
|
|
assert (
|
|
_normalize_rmbg_segment_url(
|
|
"https://example.com/api/v1/segment/fashion/",
|
|
"http://unused/v1/segment/fashion",
|
|
"segment/fashion",
|
|
)
|
|
== "https://example.com/api/v1/segment/fashion"
|
|
)
|
|
assert _health_url("https://example.com/v1/rmbg") == "https://example.com/health"
|
|
assert _health_url("https://example.com/api/v1/segment/body") == "https://example.com/api/health"
|
|
|
|
|
|
def test_rmbg_preset_generates_all_node_endpoints(monkeypatch):
|
|
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
cases = (
|
|
(ZhiYiRMBGNode(), "rmbg", "http://unused/v1/rmbg"),
|
|
(ZhiYiClothesSegmentNode(), "segment/clothes", "http://unused/v1/segment/clothes"),
|
|
(ZhiYiFashionSegmentNode(), "segment/fashion", "http://unused/v1/segment/fashion"),
|
|
(ZhiYiBodySegmentNode(), "segment/body", "http://unused/v1/segment/body"),
|
|
)
|
|
|
|
for node, endpoint, default_url in cases:
|
|
captured_urls = []
|
|
|
|
def fake_run_concurrent(tasks, max_concurrency, label):
|
|
captured_urls.append(tasks[0][2][2])
|
|
return [{
|
|
"image": image,
|
|
"mask": torch.zeros((1, 2, 2), dtype=torch.float32),
|
|
"mask_image": image,
|
|
"info": {},
|
|
}]
|
|
|
|
monkeypatch.setattr(node, "_run_concurrent", fake_run_concurrent)
|
|
_, _, _, info_text = node._execute(
|
|
image=image,
|
|
service_url="https://custom.example.com/v1/ignored",
|
|
endpoint=endpoint,
|
|
default_url=default_url,
|
|
health_key="test",
|
|
label="test",
|
|
process_res=512,
|
|
mask_blur=0,
|
|
mask_offset=0,
|
|
invert_output=False,
|
|
background="Alpha",
|
|
background_color="#222222",
|
|
return_image=True,
|
|
return_mask=True,
|
|
return_mask_image=False,
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
service_url_preset="10.1.0.230",
|
|
extra={},
|
|
info_extra={},
|
|
)
|
|
|
|
expected_url = f"http://10.1.0.230:8003/v1/{endpoint}"
|
|
assert captured_urls == [expected_url]
|
|
assert json.loads(info_text)["service_url"] == expected_url
|
|
|
|
|
|
def test_parse_classes_preserves_fashion_comma_names():
|
|
clothes = ZhiYiClothesSegmentNode()
|
|
fashion = ZhiYiFashionSegmentNode()
|
|
body = ZhiYiBodySegmentNode()
|
|
|
|
assert clothes._parse_classes_text("Upper-clothes, Pants", CLOTHES_CLASSES, allow_comma_split=True) == ["Upper-clothes", "Pants"]
|
|
assert fashion._parse_classes_text("shirt, blouse\npants", FASHION_CLASSES, allow_comma_split=False) == ["shirt, blouse", "pants"]
|
|
assert fashion._parse_classes_text('["top, t-shirt, sweatshirt", "bag, wallet"]', FASHION_CLASSES) == [
|
|
"top, t-shirt, sweatshirt",
|
|
"bag, wallet",
|
|
]
|
|
assert body._parse_classes_text("", BODY_CLASSES, allow_comma_split=True) == []
|
|
assert body._parse_classes_text(
|
|
'["Face", "Torso-skin"]', BODY_CLASSES, allow_comma_split=True
|
|
) == ["Face", "Torso-skin"]
|
|
assert body._parse_classes_text(
|
|
"Face\nHair\nTop-clothes", BODY_CLASSES, allow_comma_split=True
|
|
) == ["Face", "Hair", "Top-clothes"]
|
|
assert body._parse_classes_text(
|
|
"Face, Hair, Bottom-clothes", BODY_CLASSES, allow_comma_split=True
|
|
) == ["Face", "Hair", "Bottom-clothes"]
|
|
assert body._parse_classes_text(["Face", "Hair"], BODY_CLASSES, allow_comma_split=True) == ["Face", "Hair"]
|
|
|
|
with pytest.raises(RuntimeError, match="不支持的分割类别"):
|
|
fashion._parse_classes_text("shirt, blouse, pants", FASHION_CLASSES, allow_comma_split=False)
|
|
|
|
|
|
def test_rmbg_posts_full_payload_and_decodes_outputs(monkeypatch):
|
|
node = ZhiYiRMBGNode()
|
|
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:8003/v1/rmbg"
|
|
assert kwargs["headers"] == {"Content-Type": "application/json"}
|
|
payload = kwargs["json"]
|
|
assert payload["image"].startswith("data:image/png;base64,")
|
|
assert payload["process_res"] == 1536
|
|
assert payload["sensitivity"] == 0.85
|
|
assert payload["mask_blur"] == 4
|
|
assert payload["mask_offset"] == -2
|
|
assert payload["invert_output"] is True
|
|
assert payload["refine_foreground"] is True
|
|
assert payload["background"] == "Color"
|
|
assert payload["background_color"] == "#ffffff"
|
|
assert payload["return_image"] is True
|
|
assert payload["return_mask"] is True
|
|
assert payload["return_mask_image"] is True
|
|
assert payload["image_format"] == "png_base64"
|
|
data = _segment_response(
|
|
request_id="rmbg-1",
|
|
classes=[],
|
|
preprocessor="RMBGPreprocessor",
|
|
processed_image=_image_data_url(color=(0, 255, 0)),
|
|
mask_image=_image_data_url(color=(255, 0, 0)),
|
|
)
|
|
data["model_name"] = "RMBG-2.0"
|
|
data["model_type"] = "background_removal"
|
|
data["model_repo"] = "1038lab/RMBG-2.0"
|
|
return DummyResponse(data=data)
|
|
|
|
monkeypatch.setattr(rmbg_module.requests, "request", fake_request)
|
|
|
|
result_image, mask, mask_image, info_text = node.remove_background(
|
|
image=images,
|
|
service_url="example.com:8003",
|
|
process_res=1536,
|
|
sensitivity=0.85,
|
|
mask_blur=4,
|
|
mask_offset=-2,
|
|
invert_output=True,
|
|
refine_foreground=True,
|
|
background="Color",
|
|
background_color="#ffffff",
|
|
return_image=True,
|
|
return_mask=True,
|
|
return_mask_image=True,
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
)
|
|
|
|
info = json.loads(info_text)
|
|
assert len(calls) == 1
|
|
assert result_image.shape == (1, 2, 2, 3)
|
|
assert mask.shape == (1, 2, 2)
|
|
assert mask_image.shape == (1, 2, 2, 3)
|
|
assert torch.equal(result_image[0, 0, 0], torch.tensor([0.0, 1.0, 0.0]))
|
|
assert torch.equal(mask, torch.tensor([[[1.0, 0.0], [0.0, 1.0]]]))
|
|
assert torch.equal(mask_image[0, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
|
|
assert info["service_url"] == "http://example.com:8003/v1/rmbg"
|
|
assert info["sensitivity"] == 0.85
|
|
assert info["refine_foreground"] is True
|
|
assert info["items"][0]["request_id"] == "rmbg-1"
|
|
assert info["items"][0]["preprocessor"] == "RMBGPreprocessor"
|
|
|
|
|
|
def test_clothes_segment_posts_classes_and_composes_from_mask_when_image_missing(monkeypatch):
|
|
node = ZhiYiClothesSegmentNode()
|
|
images = torch.ones((1, 2, 2, 3), dtype=torch.float32)
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
assert method == "POST"
|
|
assert url == "https://example.com/v1/segment/clothes"
|
|
payload = kwargs["json"]
|
|
assert payload["classes"] == ["Upper-clothes", "Pants"]
|
|
assert payload["process_res"] == 512
|
|
assert payload["return_image"] is False
|
|
assert payload["return_mask"] is True
|
|
return DummyResponse(data=_segment_response(
|
|
request_id="cloth-1",
|
|
processed_image=None,
|
|
mask=_mask_data_url(white_pixels=[(0, 0), (1, 1)]),
|
|
classes=["Upper-clothes", "Pants"],
|
|
))
|
|
|
|
monkeypatch.setattr(rmbg_module.requests, "request", fake_request)
|
|
|
|
result_image, mask, _, info_text = node.segment(
|
|
image=images,
|
|
service_url="https://example.com/v1/segment/clothes",
|
|
classes="Upper-clothes\nPants",
|
|
process_res=512,
|
|
mask_blur=0,
|
|
mask_offset=0,
|
|
invert_output=False,
|
|
background="Color",
|
|
background_color="#000000",
|
|
return_image=False,
|
|
return_mask=True,
|
|
return_mask_image=False,
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
)
|
|
|
|
info = json.loads(info_text)
|
|
assert torch.equal(mask, torch.tensor([[[1.0, 0.0], [0.0, 1.0]]]))
|
|
assert torch.equal(result_image[0, 0, 0], torch.tensor([1.0, 1.0, 1.0]))
|
|
assert torch.equal(result_image[0, 0, 1], torch.tensor([0.0, 0.0, 0.0]))
|
|
assert info["classes"] == ["Upper-clothes", "Pants"]
|
|
assert info["items"][0]["classes"] == ["Upper-clothes", "Pants"]
|
|
|
|
|
|
def test_body_health_check_and_process_res_ignored_notice(monkeypatch):
|
|
node = ZhiYiBodySegmentNode()
|
|
calls = []
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
calls.append((method, url))
|
|
if method == "GET":
|
|
return DummyResponse(data={"status": "ok", "body_segment": {"loaded": True}})
|
|
payload = kwargs["json"]
|
|
assert payload["classes"] == ["Face", "Torso-skin"]
|
|
assert payload["process_res"] == 2048
|
|
return DummyResponse(data=_segment_response(
|
|
request_id="body-1",
|
|
classes=["Face", "Torso-skin"],
|
|
preprocessor="BodySegmentPreprocessor",
|
|
))
|
|
|
|
monkeypatch.setattr(rmbg_module.requests, "request", fake_request)
|
|
|
|
_, _, _, info_text = node.segment(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://example.com/api/v1/segment/body",
|
|
classes='["Face", "Torso-skin"]',
|
|
process_res=2048,
|
|
mask_blur=0,
|
|
mask_offset=0,
|
|
invert_output=False,
|
|
background="Alpha",
|
|
background_color="#222222",
|
|
return_image=True,
|
|
return_mask=True,
|
|
return_mask_image=False,
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=True,
|
|
node_switch=0,
|
|
)
|
|
|
|
info = json.loads(info_text)
|
|
assert calls[0] == ("GET", "https://example.com/api/health")
|
|
assert calls[1] == ("POST", "https://example.com/api/v1/segment/body")
|
|
assert info["process_res_ignored"] is True
|
|
assert info["items"][0]["preprocessor"] == "BodySegmentPreprocessor"
|
|
|
|
|
|
def test_node_switch_returns_original_and_empty_mask_without_request(monkeypatch):
|
|
node = ZhiYiFashionSegmentNode()
|
|
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(rmbg_module.requests, "request", fake_request)
|
|
|
|
result_image, mask, mask_image, info_text = node.segment(
|
|
image=image,
|
|
service_url="https://example.com/v1/segment/fashion",
|
|
classes="shirt, blouse",
|
|
process_res=512,
|
|
mask_blur=0,
|
|
mask_offset=0,
|
|
invert_output=False,
|
|
background="Alpha",
|
|
background_color="#222222",
|
|
return_image=True,
|
|
return_mask=True,
|
|
return_mask_image=False,
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=True,
|
|
node_switch=1,
|
|
)
|
|
|
|
assert torch.equal(result_image, image)
|
|
assert torch.equal(mask, torch.zeros((1, 2, 2), dtype=torch.float32))
|
|
assert torch.equal(mask_image, torch.zeros((1, 2, 2, 3), dtype=torch.float32))
|
|
assert json.loads(info_text)["skipped"] is True
|
|
|
|
|
|
def test_post_process_uses_service_error_message():
|
|
node = ZhiYiClothesSegmentNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
return DummyResponse(
|
|
status_code=422,
|
|
data={"error": {"code": "INVALID_CLASSES", "message": "bad class", "request_id": "req-err"}},
|
|
)
|
|
|
|
original_request = rmbg_module.requests.request
|
|
rmbg_module.requests.request = fake_request
|
|
try:
|
|
with pytest.raises(RuntimeError, match="INVALID_CLASSES: bad class: request_id=req-err"):
|
|
node._post_process("https://example.com/v1/segment/clothes", {"image": "x"}, timeout=30, label="知衣衣物语义分割")
|
|
finally:
|
|
rmbg_module.requests.request = original_request
|
|
|
|
|
|
def test_timeout_is_normalized(monkeypatch):
|
|
node = ZhiYiRMBGNode()
|
|
|
|
def fake_request(method, url, timeout, **kwargs):
|
|
raise requests.exceptions.Timeout("read timed out")
|
|
|
|
monkeypatch.setattr(rmbg_module.requests, "request", fake_request)
|
|
|
|
with pytest.raises(RuntimeError, match="TIMEOUT"):
|
|
node.remove_background(
|
|
image=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
|
service_url="https://example.com/v1/rmbg",
|
|
process_res=1024,
|
|
sensitivity=1.0,
|
|
mask_blur=0,
|
|
mask_offset=0,
|
|
invert_output=False,
|
|
refine_foreground=False,
|
|
background="Alpha",
|
|
background_color="#222222",
|
|
return_image=True,
|
|
return_mask=True,
|
|
return_mask_image=False,
|
|
max_concurrency=1,
|
|
timeout=30,
|
|
health_check=False,
|
|
node_switch=0,
|
|
)
|