241 lines
9.2 KiB
Python
241 lines
9.2 KiB
Python
"""
|
|
ComfyUI Image Annotator Node - Backend Logic (V8.2)
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
class ImageAnnotator:
|
|
def __init__(self):
|
|
# Print current file path for debugging which version is loaded
|
|
print(f"### [ImageAnnotator] Node source loaded from: {os.path.abspath(__file__)}")
|
|
self._font_cache = {} # Cache fonts by size for performance
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": { "image": ("IMAGE",), },
|
|
"optional": {
|
|
"mark_data": ("STRING", { "default": "", "multiline": True, }),
|
|
},
|
|
"hidden": { "unique_id": "UNIQUE_ID", }
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image_out", "mark_data")
|
|
FUNCTION = "annotate"
|
|
CATEGORY = "image/annotation"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("nan")
|
|
|
|
def annotate(self, image, mark_data="", unique_id=None):
|
|
print(f"### [ImageAnnotator] Processing annotation data, node path: {os.path.abspath(__file__)}")
|
|
|
|
if isinstance(image, torch.Tensor):
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(img_np)
|
|
else:
|
|
pil_image = Image.fromarray(image)
|
|
|
|
marks = []
|
|
output_mark_data = mark_data
|
|
|
|
if mark_data and mark_data.strip():
|
|
try:
|
|
data = json.loads(mark_data)
|
|
marks = data.get("marks", [])
|
|
except Exception as e:
|
|
print(f"### [ImageAnnotator] JSON parse failed: {e}")
|
|
marks = []
|
|
|
|
if not output_mark_data or not output_mark_data.strip():
|
|
output_mark_data = json.dumps({
|
|
"version": "1.0",
|
|
"image_size": [pil_image.width, pil_image.height],
|
|
"marks": []
|
|
}, ensure_ascii=False, indent=2)
|
|
|
|
if marks:
|
|
pil_image = self._render_marks(pil_image, marks)
|
|
|
|
img_np = np.array(pil_image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_np).unsqueeze(0)
|
|
|
|
return (img_tensor, output_mark_data)
|
|
|
|
def _render_marks(self, image: Image.Image, marks: list) -> Image.Image:
|
|
if image.mode != "RGBA":
|
|
image = image.convert("RGBA")
|
|
|
|
overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
|
|
draw = ImageDraw.Draw(overlay)
|
|
|
|
for idx, mark in enumerate(marks, start=1):
|
|
if not mark.get("visible", True):
|
|
continue
|
|
|
|
mark_type = mark.get("type", "point")
|
|
points = mark.get("points", [])
|
|
style = mark.get("style", {})
|
|
label = mark.get("label", "")
|
|
|
|
# Use safer defensive logic to prevent None causing lstrip crash
|
|
stroke_color = style.get("stroke_color") or "#22C55E"
|
|
stroke_width = style.get("stroke_width", 2)
|
|
fill_color = style.get("fill_color") or stroke_color
|
|
fill_alpha = style.get("fill_alpha", 0.35)
|
|
font_size = style.get("font_size", 12)
|
|
|
|
point_size = style.get("point_size", 12)
|
|
show_stroke = style.get("show_stroke", True)
|
|
show_fill = style.get("show_fill", True)
|
|
|
|
stroke_rgba = self._hex_to_rgba(stroke_color, 255)
|
|
fill_rgba = self._hex_to_rgba(fill_color, int(fill_alpha * 255))
|
|
|
|
if mark_type == "point":
|
|
if len(points) < 1:
|
|
continue
|
|
self._draw_point(draw, points, stroke_rgba, fill_rgba, stroke_width, str(idx), point_size, show_fill, show_stroke)
|
|
lx, ly = points[0]
|
|
elif mark_type == "rect":
|
|
if len(points) < 2:
|
|
continue
|
|
self._draw_rect(draw, points, stroke_rgba, fill_rgba, stroke_width, show_fill, show_stroke)
|
|
x = min(points[0][0], points[1][0])
|
|
y = min(points[0][1], points[1][1])
|
|
w = abs(points[1][0] - points[0][0])
|
|
h = abs(points[1][1] - points[0][1])
|
|
lx, ly = x + w / 2, y + h / 2
|
|
elif mark_type == "polygon":
|
|
if len(points) < 3:
|
|
continue
|
|
self._draw_polygon(draw, points, stroke_rgba, fill_rgba, stroke_width, show_fill, show_stroke)
|
|
# Calculate centroid
|
|
lx = sum(p[0] for p in points) / len(points)
|
|
ly = sum(p[1] for p in points) / len(points)
|
|
else:
|
|
continue
|
|
|
|
# Draw index number (Refined: Rectangle and Polygon use centered numbers)
|
|
if mark_type != "point":
|
|
fontSize = int(point_size * 1.2)
|
|
self._draw_centered_text(draw, lx, ly, str(idx), fontSize)
|
|
|
|
result = Image.alpha_composite(image, overlay)
|
|
return result.convert("RGB")
|
|
|
|
def _hex_to_rgba(self, hex_color, alpha: int) -> tuple:
|
|
"""Convert hex color to RGBA with defensive programming"""
|
|
default = (34, 197, 94, alpha)
|
|
|
|
if not isinstance(hex_color, str) or not hex_color:
|
|
return default
|
|
|
|
hex_color = hex_color.lstrip("#")
|
|
if len(hex_color) == 3:
|
|
hex_color = "".join([c * 2 for c in hex_color])
|
|
|
|
if len(hex_color) != 6:
|
|
return default
|
|
|
|
try:
|
|
r = int(hex_color[0:2], 16)
|
|
g = int(hex_color[2:4], 16)
|
|
b = int(hex_color[4:6], 16)
|
|
return (r, g, b, alpha)
|
|
except ValueError:
|
|
return default
|
|
|
|
def _get_font(self, size):
|
|
"""Get font with caching for performance"""
|
|
if size in self._font_cache:
|
|
return self._font_cache[size]
|
|
|
|
font_candidates = [
|
|
# macOS
|
|
"/System/Library/Fonts/PingFang.ttc",
|
|
"/System/Library/Fonts/Helvetica.ttc",
|
|
"/Library/Fonts/Arial.ttf",
|
|
# Windows
|
|
"C:/Windows/Fonts/msyh.ttc",
|
|
"C:/Windows/Fonts/simhei.ttf",
|
|
"C:/Windows/Fonts/arial.ttf",
|
|
# Linux
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
|
"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
|
|
"/usr/share/fonts/truetype/freefont/FreeSans.ttf",
|
|
]
|
|
|
|
font = None
|
|
for font_path in font_candidates:
|
|
try:
|
|
if os.path.exists(font_path):
|
|
font = ImageFont.truetype(font_path, size)
|
|
break
|
|
except:
|
|
continue
|
|
|
|
if font is None:
|
|
try:
|
|
font = ImageFont.load_default()
|
|
except:
|
|
font = None
|
|
|
|
self._font_cache[size] = font
|
|
return font
|
|
|
|
def _draw_point(self, draw: ImageDraw.Draw, points, stroke_rgba, fill_rgba, stroke_width, text, point_size, show_fill=True, show_stroke=True):
|
|
if not points: return
|
|
x, y = points[0]
|
|
radius = point_size
|
|
if show_fill:
|
|
draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=fill_rgba)
|
|
if show_stroke:
|
|
draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline=stroke_rgba, width=stroke_width)
|
|
|
|
# Match frontend: draw index badge inside point (using stroke color)
|
|
badge_r = max(8, int(radius * 0.85))
|
|
draw.ellipse([x - badge_r, y - badge_r, x + badge_r, y + badge_r], fill=stroke_rgba)
|
|
draw.ellipse(
|
|
[x - badge_r, y - badge_r, x + badge_r, y + badge_r],
|
|
outline=(255, 255, 255, 255),
|
|
width=max(1, int(stroke_width))
|
|
)
|
|
self._draw_centered_text(draw, x, y, text, max(9, int(badge_r * 1.1)))
|
|
|
|
def _draw_centered_text(self, draw: ImageDraw.Draw, x, y, text, radius):
|
|
"""Draw centered text with cached font"""
|
|
fs = max(8, int(radius * 1.2))
|
|
font = self._get_font(fs)
|
|
|
|
text_color = (255, 255, 255, 255)
|
|
bbox = draw.textbbox((0, 0), text, font=font)
|
|
tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
|
draw.text((x - tw/2, y - th/2 - 1), text, font=font, fill=text_color)
|
|
|
|
def _draw_rect(self, draw: ImageDraw.Draw, points, stroke_rgba, fill_rgba, stroke_width, show_fill=True, show_stroke=True):
|
|
if len(points) < 2: return
|
|
p1, p2 = points[0], points[1]
|
|
l, r = min(p1[0], p2[0]), max(p1[0], p2[0])
|
|
t, b = min(p1[1], p2[1]), max(p1[1], p2[1])
|
|
f_color = fill_rgba if show_fill else None
|
|
o_color = stroke_rgba if show_stroke else None
|
|
draw.rectangle([l, t, r, b], fill=f_color, outline=o_color, width=stroke_width)
|
|
|
|
def _draw_polygon(self, draw: ImageDraw.Draw, points, stroke_rgba, fill_rgba, stroke_width, show_fill=True, show_stroke=True):
|
|
if len(points) < 3: return
|
|
pts = [(p[0], p[1]) for p in points]
|
|
f_color = fill_rgba if show_fill else None
|
|
o_color = stroke_rgba if show_stroke else None
|
|
draw.polygon(pts, fill=f_color, outline=o_color)
|
|
if show_stroke:
|
|
for i in range(len(pts)):
|
|
draw.line([pts[i], pts[(i+1)%len(pts)]], fill=stroke_rgba, width=stroke_width)
|