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
115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
"""Object detection for anime images — faces, hands, persons, nudity, text, and YOLO-tagged objects."""
|
|
|
|
import importlib
|
|
import json
|
|
from collections import Counter
|
|
|
|
from comfy_api.latest import io
|
|
|
|
from .._shared.tensor import comfy_to_pil
|
|
from .._shared.formatting import label_display
|
|
|
|
|
|
class ImgUtilsDetect(io.ComfyNode):
|
|
"""Detect faces, hands, bodies, nudity, text, and Booru objects in anime images."""
|
|
|
|
DROPDOWN_OPTIONS = [
|
|
"Detect Faces",
|
|
"Detect Hands",
|
|
"Detect Heads",
|
|
"Detect Eyes",
|
|
"Detect Person",
|
|
"Detect Halfbody",
|
|
"Detect with NudeNet",
|
|
"Detect Censors",
|
|
"Detect Text with OCR",
|
|
"Detect with Booru YOLO",
|
|
]
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="ImgUtilsDetect",
|
|
display_name="Imgutils Detect",
|
|
category="imgutils/detect",
|
|
description=(
|
|
"Detect faces, hands, heads, eyes, persons, half-bodies, "
|
|
"text regions, nude areas, censors, and Booru-tagged objects in anime images."
|
|
),
|
|
search_aliases=[
|
|
"detect", "detection", "face", "hand", "head", "eye",
|
|
"person", "nudenet", "censor", "booru", "yolo", "bbox",
|
|
"text", "ocr",
|
|
],
|
|
inputs=[
|
|
io.Image.Input("image", tooltip="Input image for object detection."),
|
|
io.Combo.Input(
|
|
"mode",
|
|
options=cls.DROPDOWN_OPTIONS,
|
|
default="Detect Faces",
|
|
tooltip="What to detect — faces, hands, heads, eyes, persons, nudity, censors, text, or Booru YOLO objects.",
|
|
),
|
|
io.Float.Input(
|
|
"confidence",
|
|
default=0.5,
|
|
min=0.0,
|
|
max=1.0,
|
|
step=0.05,
|
|
tooltip="Minimum confidence threshold. Detections below this are filtered out.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.String.Output(display_name="detections"),
|
|
io.String.Output(display_name="json"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, image, mode, confidence) -> io.NodeOutput:
|
|
pil_image = comfy_to_pil(image)
|
|
results = _run_detection(pil_image, mode, confidence)
|
|
|
|
if not results:
|
|
return io.NodeOutput(
|
|
f"No detections at confidence >= {confidence:.2f}",
|
|
"[]",
|
|
)
|
|
|
|
counts = Counter(label for bbox, label, conf in results)
|
|
detections_str = ", ".join(
|
|
f"{label_display(label)}: {count}" for label, count in sorted(counts.items())
|
|
)
|
|
bboxes_json = json.dumps([
|
|
{"bbox": [bbox[0], bbox[1], bbox[2], bbox[3]], "label": label_display(label), "score": conf}
|
|
for bbox, label, conf in results
|
|
], ensure_ascii=False)
|
|
return io.NodeOutput(detections_str, bboxes_json)
|
|
|
|
|
|
def _run_detection(pil_image, mode: str, confidence: float) -> list[tuple]:
|
|
"""Run a detection function and return raw results."""
|
|
detect_map = {
|
|
"Detect Faces": "imgutils.detect.detect_faces",
|
|
"Detect Hands": "imgutils.detect.detect_hands",
|
|
"Detect Heads": "imgutils.detect.detect_heads",
|
|
"Detect Eyes": "imgutils.detect.detect_eyes",
|
|
"Detect Person": "imgutils.detect.detect_person",
|
|
"Detect Halfbody": "imgutils.detect.detect_halfbody",
|
|
"Detect with NudeNet": "imgutils.detect.detect_with_nudenet",
|
|
"Detect Censors": "imgutils.detect.detect_censors",
|
|
"Detect Text with OCR": "imgutils.ocr.detect_text_with_ocr",
|
|
"Detect with Booru YOLO": "imgutils.detect.detect_with_booru_yolo",
|
|
}
|
|
|
|
module_name, func_name = detect_map[mode].rsplit(".", 1)
|
|
mod = importlib.import_module(module_name)
|
|
raw = getattr(mod, func_name)(pil_image)
|
|
|
|
filtered = []
|
|
for item in raw:
|
|
bbox, label, conf = item
|
|
if conf >= confidence:
|
|
filtered.append(item)
|
|
|
|
return filtered
|