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
64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
"""OCR text extraction from anime images via PaddleOCR."""
|
|
|
|
import json
|
|
|
|
from comfy_api.latest import io
|
|
|
|
from .._shared.tensor import comfy_to_pil
|
|
|
|
|
|
class ImgUtilsOCR(io.ComfyNode):
|
|
"""Extract text from anime images using PaddleOCR."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="ImgUtilsOCR",
|
|
display_name="Imgutils OCR",
|
|
category="imgutils/ocr",
|
|
description=(
|
|
"Extract text from images using PaddleOCR. "
|
|
"Returns recognized text strings and text with confidence scores."
|
|
),
|
|
search_aliases=[
|
|
"ocr", "text", "recognize", "read", "paddleocr", "extract",
|
|
],
|
|
inputs=[
|
|
io.Image.Input("image", tooltip="Input image for text extraction."),
|
|
io.Float.Input(
|
|
"heat_threshold", default=0.3, min=0.0, max=1.0, step=0.05,
|
|
tooltip="Heat map threshold for text region detection. Lower = more regions.",
|
|
),
|
|
io.Float.Input(
|
|
"box_threshold", default=0.7, min=0.0, max=1.0, step=0.05,
|
|
tooltip="Box threshold for text detection. Lower = more candidate boxes.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.String.Output(display_name="detections"),
|
|
io.String.Output(display_name="json"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, image, heat_threshold, box_threshold) -> io.NodeOutput:
|
|
from imgutils.ocr import ocr
|
|
|
|
pil_image = comfy_to_pil(image)
|
|
results = ocr(
|
|
pil_image,
|
|
heat_threshold=heat_threshold,
|
|
box_threshold=box_threshold,
|
|
)
|
|
|
|
if not results:
|
|
return io.NodeOutput("No text detected.", "[]")
|
|
|
|
text_strs = [text for bbox, text, conf in results]
|
|
text_only = ", ".join(f'"{t}"' for t in text_strs)
|
|
bboxes_json = json.dumps([
|
|
{"bbox": [bbox[0], bbox[1], bbox[2], bbox[3]], "label": text, "score": conf}
|
|
for bbox, text, conf in results
|
|
], ensure_ascii=False)
|
|
return io.NodeOutput(text_only, bboxes_json)
|