xiaden
315666c7b6
- Split flat node modules into categorized subpackages (tagging, detect, edge, pose, segment, restore, compare, judge, transform, utility) - Extract shared internals into _shared/ (bases, tensor, formatting) - Centralize node registration via register_nodes() helper - Clean up: remove dead is_scored return, unnecessary hasattr guard, inline imports
50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
"""Anime character segmentation (background removal) using isnetis."""
|
|
|
|
from comfy_api.latest import io, ui
|
|
|
|
from .._shared.tensor import comfy_to_pil, pil_to_comfy
|
|
|
|
|
|
class ImgUtilsSegment(io.ComfyNode):
|
|
"""Segment anime character from background — outputs mask + RGBA."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="ImgUtilsSegment",
|
|
display_name="Imgutils Segment",
|
|
category="imgutils/segment",
|
|
description=(
|
|
"Segment anime character from background using isnetis. "
|
|
"Outputs the alpha mask and an RGBA image with transparent background."
|
|
),
|
|
search_aliases=[
|
|
"segment", "isnetis", "mask", "character",
|
|
"background", "remove bg", "cutout", "alpha",
|
|
],
|
|
inputs=[
|
|
io.Image.Input("image", tooltip="Input image containing character to segment."),
|
|
io.Int.Input(
|
|
"scale", default=1024, min=256, max=2048, step=128,
|
|
tooltip="Processing resolution. Higher = more accurate, slower. Default 1024.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Mask.Output(display_name="mask"),
|
|
io.Image.Output(display_name="image"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, image, scale=1024) -> io.NodeOutput:
|
|
from imgutils.segment import segment_rgba_with_isnetis
|
|
|
|
pil_image = comfy_to_pil(image)
|
|
|
|
mask_pil, rgba_pil = segment_rgba_with_isnetis(pil_image, scale=scale)
|
|
|
|
mask_tensor = pil_to_comfy(mask_pil)
|
|
rgba_tensor = pil_to_comfy(rgba_pil)
|
|
|
|
return io.NodeOutput(mask_tensor, rgba_tensor, ui=ui.PreviewImage(rgba_tensor, cls=cls))
|