From c81663f9a295b9c41c75771d7d05f2d20bd3a32e Mon Sep 17 00:00:00 2001 From: hua zhuang <30758919@qq.com> Date: Wed, 20 May 2026 13:01:58 +0800 Subject: [PATCH] Initial commit: ComfyUI-EnglishAnnotator --- .gitignore | 24 +++ CHANGELOG.md | 25 +++ CONTRIBUTING.md | 27 +++ INSTALL_NOTES.md | 81 ++++++++ LICENSE | 21 ++ README.md | 122 ++++++++++++ __init__.py | 21 ++ cli/README.md | 82 ++++++++ cli/english_card_cli.py | 316 +++++++++++++++++++++++++++++++ example_workflows/classroom.json | 84 ++++++++ nodes/__init__.py | 0 nodes/annotator.py | 243 ++++++++++++++++++++++++ pyproject.toml | 31 +++ requirements.txt | 1 + 14 files changed, 1078 insertions(+) create mode 100644 .gitignore create mode 100644 CHANGELOG.md create mode 100644 CONTRIBUTING.md create mode 100644 INSTALL_NOTES.md create mode 100644 LICENSE create mode 100644 README.md create mode 100644 __init__.py create mode 100644 cli/README.md create mode 100644 cli/english_card_cli.py create mode 100644 example_workflows/classroom.json create mode 100644 nodes/__init__.py create mode 100644 nodes/annotator.py create mode 100644 pyproject.toml create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..04c9162 --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +build/ +dist/ + +.venv/ +venv/ +env/ +.uv/ + +.idea/ +.vscode/ +*.swp +.DS_Store + +out/ +output/ +*.tar.gz + +# CLI / test artifacts +/tmp_workflow.json +_test_*.json diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..dd5df26 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,25 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [0.1.0] - 2026-05-20 + +### Added +- Initial release of `EnglishWordAnnotator` node. +- 8-color palette for visually distinct adjacent labels. +- Label positioning: `above`, `inside_top`, `below`, with automatic fallback when out of bounds. +- Cross-platform font fallback (Linux DejaVu, macOS Helvetica, Windows Arial). +- Robust input coercion via `_coerce_bbox_list` / `_coerce_label_list`: + bboxes accept flat list, batch-wrapped list, tensors, or list of `{"bbox":..., "label":...}` dicts; + labels accept comma/newline strings, lists of strings, or lists of dicts. +- `AnyType` input matcher for `bboxes` / `labels` so the node works across Florence-2 wrapper versions. +- CLI: `cli/english_card_cli.py` — end-to-end generate-and-annotate via ComfyUI WebSocket API, + PEP 723 inline deps for `uv run` self-contained execution. +- Example API-format workflow: `example_workflows/classroom.json`. + +### Known limitations +- No collision avoidance between label boxes — overlapping bboxes can overlap labels. +- Single image at a time (batch size 1). diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..dcbf9c9 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,27 @@ +# Contributing + +Issues and PRs are welcome. + +## PR guidelines + +1. Test against the latest `kijai/ComfyUI-Florence2`. If their output schema changes, + add a new branch to `_coerce_bbox_list` / `_coerce_label_list` in `nodes/annotator.py` + — additive only, don't break existing shapes. +2. Keep the node API stable. New options go into the `optional` section of `INPUT_TYPES`, + never into `required`, so existing workflows keep loading. +3. Add a CHANGELOG entry under an `[Unreleased]` heading. +4. Run a smoke test before submitting: + ```bash + curl -s http://localhost:8188/object_info | python3 -c \ + "import json,sys; print('EnglishWordAnnotator' in json.load(sys.stdin))" + ``` + +## Local dev loop + +```bash +# Edit nodes/annotator.py, then restart ComfyUI to pick up changes. +pkill -9 -f main.py; sleep 3 +cd /path/to/ComfyUI && source .venv/bin/activate +export HF_HOME=/path/to/hf_cache +nohup python main.py --listen 0.0.0.0 --port 8188 > ~/comfy_start.log 2>&1 & +``` diff --git a/INSTALL_NOTES.md b/INSTALL_NOTES.md new file mode 100644 index 0000000..8f0f381 --- /dev/null +++ b/INSTALL_NOTES.md @@ -0,0 +1,81 @@ +# Install Notes + +Real-world installation notes — useful when things don't go smoothly. + +## Target environment + +- Ubuntu 22.04 + RTX 3090(s) + CUDA 12.6 +- Python 3.13 inside ComfyUI's `.venv` (uv-managed) +- ComfyUI at `/mnt/modelall/ComfyUI/` +- `HF_HOME=/mnt/modelall/hf_cache` (root partition tight, must redirect) + +## Steps + +```bash +cd /mnt/modelall/ComfyUI/custom_nodes + +# Florence-2 + KJNodes — required upstream nodes +git clone https://github.com/kijai/ComfyUI-Florence2.git +git clone https://github.com/kijai/ComfyUI-KJNodes.git + +# This node +git clone https://github.com/aadebuger/ComfyUI-EnglishAnnotator.git + +# Install deps in ComfyUI venv +cd /mnt/modelall/ComfyUI && source .venv/bin/activate +export HF_HOME=/mnt/modelall/hf_cache # don't forget + +uv pip install -r custom_nodes/ComfyUI-Florence2/requirements.txt +uv pip install -r custom_nodes/ComfyUI-KJNodes/requirements.txt +uv pip install -r custom_nodes/ComfyUI-EnglishAnnotator/requirements.txt +``` + +## Model pre-download (optional but recommended) + +Avoids waiting for the first workflow to download Florence-2 weights: + +```bash +export HF_HOME=/mnt/modelall/hf_cache +uv pip install huggingface_hub +python -c " +from huggingface_hub import snapshot_download +snapshot_download('microsoft/Florence-2-base') +snapshot_download('microsoft/Florence-2-large') +" +``` + +## Start ComfyUI + +```bash +pkill -9 -f main.py; sleep 3 +cd /mnt/modelall/ComfyUI && source .venv/bin/activate +export HF_HOME=/mnt/modelall/hf_cache +nohup python main.py --listen 0.0.0.0 --port 8188 > ~/comfy_start.log 2>&1 & +sleep 30 +grep -iE "IMPORT FAILED|error" ~/comfy_start.log | head -10 +``` + +## Verify nodes loaded + +```bash +curl -s http://10.8.0.12:8188/object_info | python3 -c " +import json, sys +d = json.load(sys.stdin) +for n in ['EnglishWordAnnotator','Florence2ModelLoader','Florence2Run']: + print(f' {n}: {\"✓\" if n in d else \"✗\"}') +" +``` + +All three should be ✓. + +## Known issues / fixes + +| Symptom | Fix | +|---|---| +| `Florence2ModelLoader` complains about `transformers` version | `uv pip install -U "transformers>=4.41"` — Florence-2 needs newer, but stay below 5.0 if other nodes lock to 4.x (e.g. Qwen3TTS / Sonic at 4.57.3) | +| `flash_attn` import error | Set `attention="sdpa"` on `Florence2Run` | +| Labels render as empty boxes / squares | Fonts missing — `sudo apt install fonts-dejavu-core` | +| HF download fills root partition | `export HF_HOME=/mnt/modelall/hf_cache` *before* first run | +| `BBOX` / `STRING` type mismatch between Florence-2 versions | This node uses an `AnyType` matcher for `bboxes` / `labels`, so it tolerates flat lists, batch-wrapped lists, and dicts. Dump the workflow with `--dump-workflow /tmp/wf.json` and check `history.outputs` for the actual shape. | +| `No module named pip` inside ComfyUI venv | uv venvs don't ship pip — use `uv pip install ...`, not `python -m pip` | +| `Multiple top-level packages` when packaging | `pyproject.toml` already declares `[tool.setuptools] packages = ["nodes"]` | diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..705786b --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Hua + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..abe231d --- /dev/null +++ b/README.md @@ -0,0 +1,122 @@ +# ComfyUI-EnglishAnnotator + +A ComfyUI custom node that draws colored bounding boxes and English word labels onto an image — designed for building **vocabulary learning cards**. + +Pairs naturally with [kijai/ComfyUI-Florence2](https://github.com/kijai/ComfyUI-Florence2): use Florence-2 in `od`, `dense_region_caption`, or `caption_to_phrase_grounding` mode to produce per-object bounding boxes and English labels, then feed them into this node to render a study-ready image. + +--- + +## What it does + +**Inputs:** +- An image (e.g. a scene generated by SDXL / Flux / Z-Image-Turbo) +- A list of bounding boxes `[[x1, y1, x2, y2], ...]` +- A list of English labels `["apple", "chair", ...]` + +**Output:** +- The same image with a colored rectangle around each detected object and the English word printed in a colored label box (above / inside / below the bbox) + +A built-in 8-color palette keeps adjacent labels visually distinct. The node is permissive about input shape — it accepts flat lists, batch-wrapped lists, or lists of `{"bbox":..., "label":...}` dicts, so it works across Florence-2 wrapper versions. + +--- + +## Install + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/aadebuger/ComfyUI-EnglishAnnotator.git +cd ComfyUI-EnglishAnnotator + +# inside ComfyUI's venv +uv pip install -r requirements.txt +``` + +Restart ComfyUI. The node appears under category **📚 English Annotator** as **📚 English Word Annotator**. + +See [`INSTALL_NOTES.md`](INSTALL_NOTES.md) for environment-specific notes (Ubuntu + uv, HF cache redirect, etc.). + +--- + +## Node reference + +| Input | Type | Default | Notes | +|---|---|---|---| +| `image` | IMAGE | — | The source image | +| `bboxes` | any | — | Bounding boxes (typically from Florence-2) | +| `labels` | any | — | Labels (typically from Florence-2) | +| `font_size` | INT | 36 | Label text size in pixels | +| `box_width` | INT | 4 | Bounding box outline width | +| `label_position` | enum | `above` | `above` / `inside_top` / `below`, auto-fallback when out of bounds | +| `font_path` | STRING (optional) | `""` | Explicit TTF/OTF path (overrides defaults) | +| `text_color` | STRING (optional) | `#FFFFFF` | Hex color for label text | +| `max_labels` | INT (optional) | `0` | Cap the number of objects to annotate (0 = no limit) | + +**Returns:** `IMAGE`. + +**Fonts:** the node walks a built-in candidate list — DejaVu (Linux), Helvetica/Arial (macOS), Arial (Windows). On a fresh Ubuntu box install `fonts-dejavu-core` if labels render as boxes. + +--- + +## Example workflow + +The flow is: + +``` +CheckpointLoader → KSampler → VAEDecode ──┐ + ├→ EnglishWordAnnotator → SaveImage + Florence2ModelLoader → Florence2Run ──┘ +``` + +See [`example_workflows/classroom.json`](example_workflows/classroom.json) for an API-format workflow you can submit via `/prompt`, or generate one with the CLI's `--dump-workflow` flag. + +### Florence-2 task suggestions + +| Task | When to use | +|---|---| +| `od` | Clean COCO-style category names — best for beginner-level vocab | +| `caption_to_phrase_grounding` | You supply your own word list; Florence-2 locates exactly those words | +| `dense_region_caption` | Natural-language regions like "red apple on table" | + +For a fixed vocabulary pipeline (e.g. a daily 16-word list), `caption_to_phrase_grounding` is by far the most reliable. + +--- + +## CLI + +A self-contained CLI is bundled under [`cli/`](cli/). It uses ComfyUI's WebSocket API to submit a workflow end-to-end and download the annotated image: + +```bash +uv run cli/english_card_cli.py \ + "flat vector illustration of a tidy classroom: desk, chair, blackboard, books, pencil, backpack, clock, globe, computer" \ + --ip 10.8.0.12 --steps 25 +``` + +With a fixed vocabulary (uses `caption_to_phrase_grounding`): + +```bash +uv run cli/english_card_cli.py \ + "a sunny kitchen scene, flat vector illustration" \ + --words "apple,banana,knife,cutting board,bowl,kettle,faucet,refrigerator" \ + --ip 10.8.0.12 +``` + +See [`cli/README.md`](cli/README.md) for full argument reference. + +--- + +## Use case: daily vocabulary card pipeline + +This node was built for an automated daily English vocabulary card workflow: + +1. A daily JSON of 16 target words (with definitions and example sentences) +2. Scene generation via SDXL / Z-Image-Turbo +3. Florence-2 `caption_to_phrase_grounding` locks bboxes onto exactly those 16 words +4. **This node** renders the annotated image +5. VoxCPM2 TTS for per-word audio +6. ffmpeg assembles the final card video + +--- + +## License + +[MIT](LICENSE) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e22f828 --- /dev/null +++ b/__init__.py @@ -0,0 +1,21 @@ +"""ComfyUI-EnglishAnnotator + +A custom node that draws bounding boxes and English word labels onto +an image — designed for building vocabulary learning cards. + +Pairs naturally with kijai/ComfyUI-Florence2 which produces the +(bboxes, labels) outputs that this node consumes. +""" +from .nodes.annotator import EnglishWordAnnotator + +NODE_CLASS_MAPPINGS = { + "EnglishWordAnnotator": EnglishWordAnnotator, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "EnglishWordAnnotator": "📚 English Word Annotator", +} + +WEB_DIRECTORY = None + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/cli/README.md b/cli/README.md new file mode 100644 index 0000000..c9dbfa9 --- /dev/null +++ b/cli/README.md @@ -0,0 +1,82 @@ +# CLI: english_card_cli.py + +End-to-end CLI for the English vocabulary card pipeline: +**text prompt → generated scene → Florence-2 detection → annotated image**. + +Uses ComfyUI's WebSocket API. PEP 723 inline dependencies, so `uv run` sets up the venv automatically. + +## Run + +Free detection (Florence-2 picks what it sees): + +```bash +uv run english_card_cli.py \ + "flat vector illustration of a tidy classroom: desk, chair, blackboard, books, pencil, backpack, clock, globe, computer" \ + --ip 10.8.0.12 --steps 25 +``` + +Fixed vocabulary (uses `caption_to_phrase_grounding` — most reliable for daily card pipelines): + +```bash +uv run english_card_cli.py \ + "a sunny kitchen scene, flat vector illustration" \ + --words "apple,banana,knife,cutting board,bowl,kettle,faucet,refrigerator" \ + --ip 10.8.0.12 +``` + +## Arguments + +| Arg | Default | Purpose | +|---|---|---| +| `prompt` | (required) | Scene prompt | +| `--words` | — | Comma-separated target words; switches Florence-2 to `caption_to_phrase_grounding` | +| `--ckpt` | `sd_xl_base_1.0.safetensors` | Checkpoint filename in `models/checkpoints/` | +| `--width` / `--height` | 1024 / 1024 | Output dimensions | +| `--steps` | 25 | Sampler steps | +| `--cfg` | 6.0 | CFG scale | +| `--seed` | (timestamp) | Sampler seed | +| `--florence2-model` | `microsoft/Florence-2-large` | HF model id | +| `--florence2-precision` | `fp16` | `fp16` / `bf16` / `fp32` | +| `--font-size` | 36 | Label text size | +| `--box-width` | 4 | Box outline width | +| `--ip` / `--port` | `127.0.0.1` / `8188` | ComfyUI server | +| `--server` | — | Full `host:port` (overrides `--ip`/`--port`) | +| `--out` | `./out` | Output directory | +| `--dump-workflow` | — | Save generated API workflow JSON for debugging | +| `--timeout` | 600 | Seconds before giving up | + +## Debugging + +Dump the workflow JSON, then inspect or replay: + +```bash +./english_card_cli.py "scene prompt" --dump-workflow /tmp/wf.json + +# 1. Load /tmp/wf.json in the ComfyUI UI (Workflow → Load (API Format)) to inspect +# 2. Or POST it directly: +curl -X POST -H "Content-Type: application/json" \ + -d "$(jq -n --slurpfile p /tmp/wf.json '{prompt: $p[0]}')" \ + http://10.8.0.12:8188/prompt +``` + +Inspect the latest history (to see actual node output shapes): + +```bash +curl -s http://10.8.0.12:8188/history | python3 -c " +import json, sys +h = json.load(sys.stdin) +last = sorted(h)[-1] +print(json.dumps(h[last]['outputs'], indent=2, default=str)[:1500]) +" +``` + +## Common adjustments + +- **BBOX index mismatch** — if Florence-2 returns bboxes at a different output index, edit the + `annotate` node in `build_workflow()`: + ```python + "bboxes": ["f2_run", 2], # try 1, 2, or 3 depending on Florence2 version + "labels": ["f2_run", 3], + ``` +- **Z-Image-Turbo** — set `--steps 4 --cfg 1.0` and use the turbo checkpoint name. +- **Flux** — adjust `sampler_name` / `scheduler` in `build_workflow()` accordingly. diff --git a/cli/english_card_cli.py b/cli/english_card_cli.py new file mode 100644 index 0000000..277a2aa --- /dev/null +++ b/cli/english_card_cli.py @@ -0,0 +1,316 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "websocket-client>=1.7", +# ] +# /// +"""English vocabulary card CLI. + +End-to-end: prompt -> SDXL/Z-Image -> Florence-2 detection -> annotator -> save. + +Examples +-------- + uv run english_card_cli.py \\ + "flat vector illustration of a tidy classroom: desk, chair, blackboard, books, pencil, backpack, clock, globe, computer" \\ + --ip 10.8.0.12 --steps 25 + + uv run english_card_cli.py \\ + "a sunny kitchen scene, flat vector illustration" \\ + --words "apple,banana,knife,cutting board,bowl,kettle,faucet,refrigerator" \\ + --ip 10.8.0.12 +""" +import argparse +import json +import sys +import time +import urllib.parse +import urllib.request +import uuid +from pathlib import Path + +import websocket # websocket-client + + +# --------------------------------------------------------------------------- +# ComfyClient — minimal HTTP + WebSocket client for the ComfyUI server. +# --------------------------------------------------------------------------- +class ComfyClient: + def __init__(self, server: str = "127.0.0.1:8188"): + self.server = server + self.client_id = str(uuid.uuid4()) + + def submit(self, workflow: dict) -> str: + body = json.dumps( + {"prompt": workflow, "client_id": self.client_id} + ).encode() + req = urllib.request.Request( + f"http://{self.server}/prompt", + data=body, + headers={"Content-Type": "application/json"}, + ) + return json.loads(urllib.request.urlopen(req).read())["prompt_id"] + + def history(self, pid: str) -> dict: + r = urllib.request.urlopen(f"http://{self.server}/history/{pid}") + return json.loads(r.read()).get(pid, {}) + + def view(self, filename: str, subfolder: str, type_: str) -> bytes: + q = urllib.parse.urlencode( + {"filename": filename, "subfolder": subfolder, "type": type_} + ) + return urllib.request.urlopen(f"http://{self.server}/view?{q}").read() + + def run(self, workflow: dict, timeout_s: float = 600.0) -> dict: + ws = websocket.WebSocket() + ws.connect(f"ws://{self.server}/ws?clientId={self.client_id}") + pid = self.submit(workflow) + start, last = time.time(), None + try: + while time.time() - start < timeout_s: + msg = ws.recv() + if not isinstance(msg, str): + continue + m = json.loads(msg) + t, d = m.get("type"), m.get("data", {}) + + if t == "progress": + cur, tot = d.get("value", 0), d.get("max", 1) + print(f"\r ▸ {cur:>4}/{tot}", end="", flush=True) + + elif t == "executing": + if d.get("prompt_id") == pid and d.get("node") is None: + print() # newline after progress + return self.history(pid) + node = d.get("node") + if node and node != last: + last = node + print(f"\r ▸ node {node:<22}", end="", flush=True) + + elif t == "execution_error": + raise RuntimeError( + f"Node {d.get('node_id')}: {d.get('exception_message')}" + ) + finally: + ws.close() + raise TimeoutError(f"timeout after {timeout_s}s") + + +# --------------------------------------------------------------------------- +# Workflow builder — returns API-format dict. +# --------------------------------------------------------------------------- +def build_workflow( + prompt: str, + *, + ckpt: str, + width: int, + height: int, + steps: int, + cfg: float, + seed: int, + target_words: list | None = None, + font_size: int = 36, + box_width: int = 4, + florence2_model: str = "microsoft/Florence-2-large", + florence2_precision: str = "fp16", +) -> dict: + """SDXL/Z-Image -> Florence-2 -> Annotator -> Save.""" + if target_words: + f2_task = "caption_to_phrase_grounding" + f2_text = ", ".join(target_words) + else: + f2_task = "dense_region_caption" + f2_text = "" + + return { + "loader": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": ckpt}, + }, + "pos": { + "class_type": "CLIPTextEncode", + "inputs": {"text": prompt, "clip": ["loader", 1]}, + }, + "neg": { + "class_type": "CLIPTextEncode", + "inputs": { + "text": "blurry, low quality, watermark, text overlay", + "clip": ["loader", 1], + }, + }, + "latent": { + "class_type": "EmptySD3LatentImage", + "inputs": {"width": width, "height": height, "batch_size": 1}, + }, + "sampler": { + "class_type": "KSampler", + "inputs": { + "model": ["loader", 0], + "seed": seed, + "steps": steps, + "cfg": cfg, + "sampler_name": "euler", + "scheduler": "simple", + "denoise": 1.0, + "positive": ["pos", 0], + "negative": ["neg", 0], + "latent_image": ["latent", 0], + }, + }, + "decode": { + "class_type": "VAEDecode", + "inputs": {"samples": ["sampler", 0], "vae": ["loader", 2]}, + }, + "f2_load": { + "class_type": "Florence2ModelLoader", + "inputs": { + "model": florence2_model, + "precision": florence2_precision, + "attention": "sdpa", + }, + }, + "f2_run": { + "class_type": "Florence2Run", + "inputs": { + "image": ["decode", 0], + "florence2_model": ["f2_load", 0], + "task": f2_task, + "text_input": f2_text, + "fill_mask": False, + "keep_model_loaded": True, + "max_new_tokens": 1024, + "num_beams": 3, + "do_sample": False, + "output_mask_select": "", + }, + }, + "annotate": { + "class_type": "EnglishWordAnnotator", + "inputs": { + "image": ["decode", 0], + # Florence2Run outputs vary across versions; the annotator + # accepts any shape. Adjust indices if your version differs. + "bboxes": ["f2_run", 2], + "labels": ["f2_run", 3], + "font_size": font_size, + "box_width": box_width, + "label_position": "above", + }, + }, + "save": { + "class_type": "SaveImage", + "inputs": { + "images": ["annotate", 0], + "filename_prefix": "english_card", + }, + }, + } + + +# --------------------------------------------------------------------------- +# Entry point. +# --------------------------------------------------------------------------- +def main(): + ap = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, + description=__doc__, + ) + ap.add_argument("prompt", help="scene prompt") + ap.add_argument( + "--words", + help="comma-separated target words; uses Florence-2 phrase grounding", + ) + + # generation params + ap.add_argument("--ckpt", default="sd_xl_base_1.0.safetensors") + ap.add_argument("--width", type=int, default=1024) + ap.add_argument("--height", type=int, default=1024) + ap.add_argument("--steps", type=int, default=25) + ap.add_argument("--cfg", type=float, default=6.0) + ap.add_argument( + "--seed", type=int, default=int(time.time()) & 0xFFFFFFFF + ) + + # florence-2 + ap.add_argument( + "--florence2-model", default="microsoft/Florence-2-large" + ) + ap.add_argument( + "--florence2-precision", default="fp16", + choices=["fp16", "bf16", "fp32"], + ) + + # annotation + ap.add_argument("--font-size", type=int, default=36) + ap.add_argument("--box-width", type=int, default=4) + + # server + ap.add_argument("--ip", default="127.0.0.1") + ap.add_argument("--port", type=int, default=8188) + ap.add_argument("--server", default=None, help="full host:port") + + # output / debug + ap.add_argument("--out", type=Path, default=Path("./out")) + ap.add_argument( + "--dump-workflow", + type=Path, + default=None, + help="save generated workflow JSON (debug)", + ) + ap.add_argument("--timeout", type=float, default=600) + + args = ap.parse_args() + server = args.server or f"{args.ip}:{args.port}" + args.out.mkdir(parents=True, exist_ok=True) + + target_words = ( + [w.strip() for w in args.words.split(",") if w.strip()] + if args.words + else None + ) + + wf = build_workflow( + args.prompt, + ckpt=args.ckpt, + width=args.width, + height=args.height, + steps=args.steps, + cfg=args.cfg, + seed=args.seed, + target_words=target_words, + font_size=args.font_size, + box_width=args.box_width, + florence2_model=args.florence2_model, + florence2_precision=args.florence2_precision, + ) + + if args.dump_workflow: + args.dump_workflow.write_text(json.dumps(wf, indent=2)) + print(f"workflow dumped -> {args.dump_workflow}") + + print(f"server: {server} seed: {args.seed}") + if target_words: + print(f"target words: {', '.join(target_words)}") + + client = ComfyClient(server) + hist = client.run(wf, timeout_s=args.timeout) + + imgs = hist.get("outputs", {}).get("save", {}).get("images", []) + if not imgs: + print("no images in history. outputs:") + print(json.dumps(hist.get("outputs", {}), indent=2, default=str)[:1200]) + sys.exit(1) + + for img in imgs: + data = client.view( + img["filename"], + img.get("subfolder", ""), + img.get("type", "output"), + ) + dst = args.out / img["filename"] + dst.write_bytes(data) + print(f"saved -> {dst} ({len(data) // 1024} KB)") + + +if __name__ == "__main__": + main() diff --git a/example_workflows/classroom.json b/example_workflows/classroom.json new file mode 100644 index 0000000..71ab45c --- /dev/null +++ b/example_workflows/classroom.json @@ -0,0 +1,84 @@ +{ + "loader": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "sd_xl_base_1.0.safetensors" } + }, + "pos": { + "class_type": "CLIPTextEncode", + "inputs": { + "text": "flat vector illustration of a tidy classroom: wooden desk, chair, blackboard, world map, books, pencil, eraser, backpack, clock, globe, computer, calendar, isometric, bright colors, every object clearly separated, soft daylight", + "clip": ["loader", 1] + } + }, + "neg": { + "class_type": "CLIPTextEncode", + "inputs": { + "text": "blurry, low quality, watermark, text overlay", + "clip": ["loader", 1] + } + }, + "latent": { + "class_type": "EmptySD3LatentImage", + "inputs": { "width": 1024, "height": 1024, "batch_size": 1 } + }, + "sampler": { + "class_type": "KSampler", + "inputs": { + "model": ["loader", 0], + "seed": 42, + "steps": 25, + "cfg": 6.0, + "sampler_name": "euler", + "scheduler": "simple", + "denoise": 1.0, + "positive": ["pos", 0], + "negative": ["neg", 0], + "latent_image": ["latent", 0] + } + }, + "decode": { + "class_type": "VAEDecode", + "inputs": { "samples": ["sampler", 0], "vae": ["loader", 2] } + }, + "f2_load": { + "class_type": "Florence2ModelLoader", + "inputs": { + "model": "microsoft/Florence-2-large", + "precision": "fp16", + "attention": "sdpa" + } + }, + "f2_run": { + "class_type": "Florence2Run", + "inputs": { + "image": ["decode", 0], + "florence2_model": ["f2_load", 0], + "task": "caption_to_phrase_grounding", + "text_input": "desk, chair, blackboard, books, pencil, backpack, clock, globe, computer", + "fill_mask": false, + "keep_model_loaded": true, + "max_new_tokens": 1024, + "num_beams": 3, + "do_sample": false, + "output_mask_select": "" + } + }, + "annotate": { + "class_type": "EnglishWordAnnotator", + "inputs": { + "image": ["decode", 0], + "bboxes": ["f2_run", 2], + "labels": ["f2_run", 3], + "font_size": 36, + "box_width": 4, + "label_position": "above" + } + }, + "save": { + "class_type": "SaveImage", + "inputs": { + "images": ["annotate", 0], + "filename_prefix": "english_card" + } + } +} diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/annotator.py b/nodes/annotator.py new file mode 100644 index 0000000..5b89be0 --- /dev/null +++ b/nodes/annotator.py @@ -0,0 +1,243 @@ +"""English Word Annotator node. + +Draws colored bounding boxes + English word labels onto an image. +Intended downstream of a detector node such as kijai/ComfyUI-Florence2 +running in `od`, `dense_region_caption`, or `caption_to_phrase_grounding` +mode. + +The node is deliberately permissive about the *shape* of the bbox and +label inputs because Florence-2 wrappers have shifted their output +schema across releases. See `_coerce_bbox_list` / `_coerce_label_list`. +""" +import os +from typing import Any, List + +import numpy as np +import torch +from PIL import Image, ImageDraw, ImageFont + + +# --- Type that matches anything (rgthree / pythongosssss pattern) ----------- +class AnyType(str): + """A 'type' that always passes ComfyUI's type-equality check.""" + + def __ne__(self, other): + return False + + +ANY = AnyType("*") + + +# --- Constants -------------------------------------------------------------- +PALETTE = [ + "#FF6B6B", "#4ECDC4", "#FFD93D", "#6BCB77", + "#4D96FF", "#FF9F43", "#A66CFF", "#FF6FB5", +] + +DEFAULT_FONT_CANDIDATES = [ + # Linux (Ubuntu/Debian) + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", + # macOS + "/System/Library/Fonts/Helvetica.ttc", + "/System/Library/Fonts/Supplemental/Arial Bold.ttf", + "/Library/Fonts/Arial.ttf", + # Windows + "C:/Windows/Fonts/arialbd.ttf", + "C:/Windows/Fonts/arial.ttf", +] + + +# --- Helpers --------------------------------------------------------------- +def _load_font(font_path: str, size: int) -> ImageFont.ImageFont: + candidates = [font_path] + DEFAULT_FONT_CANDIDATES + for p in candidates: + if not p: + continue + if os.path.exists(p): + try: + return ImageFont.truetype(p, size) + except Exception: + continue + return ImageFont.load_default() + + +def _coerce_bbox_list(bboxes: Any) -> List[List[float]]: + """Normalize many possible bbox shapes to [[x1,y1,x2,y2], ...]. + + Accepted inputs: + - None / empty + - torch.Tensor of shape (N, 4) or (1, N, 4) + - [[x1,y1,x2,y2], ...] + - [[[x1,y1,x2,y2], ...]] (batch-wrapped) + - [{"bbox": [...], "label": "..."}, ...] (older Florence2) + - [[{"bbox": [...], "label": "..."}, ...]] (batch-wrapped dicts) + """ + if bboxes is None: + return [] + + if torch.is_tensor(bboxes): + bboxes = bboxes.tolist() + + if not isinstance(bboxes, (list, tuple)): + return [] + + # Unwrap a single-batch outer list. + if ( + len(bboxes) == 1 + and isinstance(bboxes[0], list) + and bboxes[0] + and isinstance(bboxes[0][0], (list, tuple, dict)) + ): + bboxes = bboxes[0] + + out: List[List[float]] = [] + for b in bboxes: + if isinstance(b, dict) and "bbox" in b: + out.append(list(b["bbox"])) + elif isinstance(b, (list, tuple)) and len(b) == 4: + out.append(list(b)) + return out + + +def _coerce_label_list(labels: Any) -> List[str]: + """Normalize many possible label shapes to a flat list of strings.""" + if labels is None: + return [] + + if isinstance(labels, str): + parts = labels.replace(",", "\n").split("\n") + return [p.strip() for p in parts if p.strip()] + + if isinstance(labels, (list, tuple)): + # Unwrap batch. + if len(labels) == 1 and isinstance(labels[0], (list, tuple)): + labels = labels[0] + + out: List[str] = [] + for x in labels: + if isinstance(x, dict): + # could be {"label": "...", ...} or {"text": "...", ...} + if "label" in x: + out.append(str(x["label"])) + elif "text" in x: + out.append(str(x["text"])) + else: + out.append(str(x)) + return out + + return [str(labels)] + + +# --- Node ------------------------------------------------------------------- +class EnglishWordAnnotator: + """Draw boxes + English word labels on an image for vocabulary cards.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "bboxes": (ANY,), + "labels": (ANY,), + "font_size": ("INT", {"default": 36, "min": 8, "max": 256}), + "box_width": ("INT", {"default": 4, "min": 1, "max": 20}), + "label_position": ( + ["above", "inside_top", "below"], + {"default": "above"}, + ), + }, + "optional": { + "font_path": ("STRING", {"default": ""}), + "text_color": ("STRING", {"default": "#FFFFFF"}), + "max_labels": ( + "INT", + {"default": 0, "min": 0, "max": 100, + "tooltip": "0 = no limit"}, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "annotate" + CATEGORY = "📚 English Annotator" + + # ----- main entrypoint ----- + def annotate( + self, + image, + bboxes, + labels, + font_size: int, + box_width: int, + label_position: str, + font_path: str = "", + text_color: str = "#FFFFFF", + max_labels: int = 0, + ): + # IMAGE tensor is (B, H, W, C) float in [0, 1]. + if image.ndim == 4: + arr = image[0] + else: + arr = image + arr = (arr.detach().cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8) + + pil = Image.fromarray(arr).convert("RGB") + draw = ImageDraw.Draw(pil) + font = _load_font(font_path, font_size) + + bbox_list = _coerce_bbox_list(bboxes) + label_list = _coerce_label_list(labels) + + if max_labels and max_labels > 0: + bbox_list = bbox_list[:max_labels] + label_list = label_list[:max_labels] + + img_w, img_h = pil.size + pad = max(4, font_size // 6) + + for i, bbox in enumerate(bbox_list): + try: + x1, y1, x2, y2 = [int(round(float(v))) for v in bbox] + except Exception: + continue + # Clip into image bounds. + x1, x2 = sorted([max(0, x1), min(img_w - 1, x2)]) + y1, y2 = sorted([max(0, y1), min(img_h - 1, y2)]) + if x2 - x1 < 2 or y2 - y1 < 2: + continue + + label = label_list[i] if i < len(label_list) else f"#{i + 1}" + color = PALETTE[i % len(PALETTE)] + + draw.rectangle([x1, y1, x2, y2], outline=color, width=box_width) + + tb = draw.textbbox((0, 0), label, font=font) + tw, th = tb[2] - tb[0], tb[3] - tb[1] + + if label_position == "above": + bg_y1, bg_y2 = y1 - th - pad * 2, y1 + if bg_y1 < 0: # doesn't fit above -> fall back to inside_top + bg_y1, bg_y2 = y1, y1 + th + pad * 2 + elif label_position == "below": + bg_y1, bg_y2 = y2, y2 + th + pad * 2 + if bg_y2 > img_h: + bg_y1, bg_y2 = y2 - th - pad * 2, y2 + else: # inside_top + bg_y1, bg_y2 = y1, y1 + th + pad * 2 + + bg_x1 = x1 + bg_x2 = min(img_w, x1 + tw + pad * 2) + + draw.rectangle([bg_x1, bg_y1, bg_x2, bg_y2], fill=color) + draw.text( + (bg_x1 + pad, bg_y1 + pad - tb[1]), + label, + fill=text_color, + font=font, + ) + + out = np.array(pil).astype(np.float32) / 255.0 + return (torch.from_numpy(out).unsqueeze(0),) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b65df26 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,31 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui-english-annotator" +version = "0.1.0" +description = "ComfyUI node that draws bounding boxes + English word labels for vocabulary learning cards" +readme = "README.md" +license = { text = "MIT" } +requires-python = ">=3.10" +authors = [{ name = "Hua" }] +keywords = ["comfyui", "stable-diffusion", "english", "education", "vocabulary", "florence-2"] +classifiers = [ + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Topic :: Multimedia :: Graphics", + "Topic :: Education", +] +dependencies = [ + "Pillow>=10.0", + "numpy", +] + +[project.urls] +Homepage = "https://github.com/aadebuger/ComfyUI-EnglishAnnotator" +Issues = "https://github.com/aadebuger/ComfyUI-EnglishAnnotator/issues" + +[tool.setuptools] +packages = ["nodes"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..750ca02 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +Pillow>=10.0