nodes/baijunty-comfyui_image_embeddings
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

```

Browse Source 
feat(nodes): 添加GLM-OCR节点并优化图像处理功能

- 新增GLMOCRNode类，支持文本、公式和表格识别功能
- 优化process_image函数，修复字符串引号格式问题
- 改进CustomImageLoader类，支持批量图像处理和更好的URL处理
- 优化Image2Base64类，改进图像张量处理逻辑
- 添加VisionOutputEmbedding2JSON节点，支持视觉输出转JSON
- 新增ImageHashNode节点，提供图像哈希计算功能
- 优化Resize2DivisibleImage节点，改进图像尺寸调整算法
- 新增Base64ImageLoader节点，支持base64图像加载
- 统一代码格式，修复数值精度和字符串引号一致性问题
```
This commit is contained in:
bai
2026-02-11 17:14:01 +08:00
parent 5623bdff8a
commit aaefb465e2
1 changed files with 241 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/image_embeddings/nodes.py
+241 -77
View File
@@ -8,13 +8,14 @@ import os
from io import BytesIO
from torchvision import transforms
def process_image(img,name):
def process_image(img, name):
output_images = []
output_masks = []
w, h = None, None
excluded_formats = ['MPO']
for i in Image.Spin(img) if hasattr(Image, 'Spin') else [img]:
if i.mode == 'I':
excluded_formats = ["MPO"]
for i in Image.Spin(img) if hasattr(Image, "Spin") else [img]:
if i.mode == "I":
i = i.point(lambda i: i * (1 / 255))
image = i.convert("RGB")
@@ -27,14 +28,16 @@ def process_image(img,name):
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
if 'A' in i.getbands():
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
elif i.mode == 'P' and 'transparency' in i.info:
mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
if "A" in i.getbands():
mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
elif i.mode == "P" and "transparency" in i.info:
mask = (
np.array(i.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0
)
mask = 1.0 - torch.from_numpy(mask)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
output_images.append(image)
output_masks.append(mask.unsqueeze(0))
@@ -44,7 +47,8 @@ def process_image(img,name):
else:
output_image = output_images[0]
output_mask = output_masks[0]
return (output_image, output_mask,name)
return (output_image, output_mask, name)
class CustomImageLoader:
@classmethod
@@ -56,12 +60,16 @@ class CustomImageLoader:
}
CATEGORY = "image"
RETURN_TYPES = ("IMAGE", "MASK","STRING",)
RETURN_TYPES = (
"IMAGE",
"MASK",
"STRING",
)
FUNCTION = "load_image"
def load_image(self, image_path_or_url):
# 检查是否为URL
if image_path_or_url.startswith(('http://', 'https://')):
if image_path_or_url.startswith(("http://", "https://")):
return self.load_image_from_url(image_path_or_url)
else:
return self.load_image_from_path(image_path_or_url)
@@ -70,7 +78,7 @@ class CustomImageLoader:
response = requests.get(url, timeout=90)
response.raise_for_status()
img = Image.open(BytesIO(response.content))
return process_image(img,url.split('/')[-1])
return process_image(img, url.split("/")[-1])
def load_image_from_path(self, image_path):
# 检查是否为相对路径或绝对路径
@@ -78,41 +86,43 @@ class CustomImageLoader:
# 如果不是绝对路径，尝试在ComfyUI的input目录中查找
input_dir = folder_paths.get_input_directory()
image_path = os.path.join(input_dir, image_path)
# 检查是否为目录
if os.path.isdir(image_path):
# 获取目录中的所有图片文件
image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'}
image_extensions = {".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"}
image_files = []
for file in os.listdir(image_path):
file_ext = os.path.splitext(file)[1].lower()
if file_ext in image_extensions:
image_files.append(os.path.join(image_path, file))
# 按文件名排序，保证顺序一致
image_files.sort()
# 加载所有图片
all_images = []
all_masks = []
all_names = []
for img_path in image_files:
img = Image.open(img_path)
processed_img, processed_mask,_ = process_image(img,'')
processed_img, processed_mask, _ = process_image(img, "")
all_images.append(processed_img)
all_masks.append(processed_mask)
all_names.append(os.path.basename(img_path))
return (all_images, all_masks,all_names)
return (torch.cat(all_images, dim=0), torch.cat(all_masks, dim=0), all_names)
else:
# 原来的单个文件处理逻辑
img = Image.open(image_path)
return process_image(img,os.path.basename(image_path))
return process_image(img, os.path.basename(image_path))
@classmethod
def IS_CHANGED(s, image_path_or_url):
return float("NaN") if image_path_or_url.startswith(('http://', 'https://')) else image_path_or_url
return (
float("NaN")
if image_path_or_url.startswith(("http://", "https://"))
else image_path_or_url
)
class Image2Base64:
@@ -121,59 +131,70 @@ class Image2Base64:
return {
"required": {
"images": ("IMAGE",),
"names": ("STRING",), },
"names": ("STRING",),
},
}
CATEGORY = "utils"
RETURN_TYPES = ('STRING',)
RETURN_TYPES = ("STRING",)
FUNCTION = "image_to_base64"
def image_to_base64(self, images,names):
def image_to_base64(self, images, names):
"""
将输入的图像张量转换为base64编码，并按名称:编码内容的格式输出JSON字符串
"""
# 创建结果字典
result_dict = {}
file_names=[]
import base64
if isinstance(names, list):
file_names = names
else:
file_names = [names]
# 遍历图像批次
for index, image in enumerate(images):
# 确保处理的是 [H, W, C] 格式
if len(image.shape) == 4 and image.shape[0] == 1:
image = image.squeeze(0)
# 转换图像张量为PIL图像
n = 255. * image.cpu().numpy()
n = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(n, 0, 255).astype(np.uint8))
# 将PIL图像转换为base64
buffer = BytesIO()
img.save(buffer, format="PNG")
img_bytes = buffer.getvalue()
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
img_base64 = base64.b64encode(img_bytes).decode("utf-8")
# 添加到结果字典
result_dict[file_names[index]] = img_base64
# 将结果字典转换为JSON字符串
json_string = json.dumps(result_dict)
return (json_string,)
class VisionOutputEmbedding2JSON:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"vision_output": ("CLIP_VISION_OUTPUT",),
"name": ("STRING",), },
"name": ("STRING",),
},
}
CATEGORY = "utils"
RETURN_TYPES = ('STRING',)
RETURN_TYPES = ("STRING",)
FUNCTION = "output_embedding_to_json"
def output_embedding_to_json(self, vision_output,name):
def output_embedding_to_json(self, vision_output, name):
# 将tensor转换为numpy数组后再进行JSON编码
image_embeds = vision_output['image_embeds']
image_embeds = vision_output["image_embeds"]
if torch.is_tensor(image_embeds):
# 如果是PyTorch张量，转换为numpy数组
image_embeds = image_embeds.squeeze().detach().cpu().numpy().tolist()
@@ -181,6 +202,8 @@ class VisionOutputEmbedding2JSON:
# 如果已经是numpy数组，转换为列表
image_embeds = image_embeds.tolist()
return (json.dumps({name: image_embeds}),)
def hex_to_signed(hex_str, bits):
"""
将十六进制字符串转换为有符号整数。
@@ -196,42 +219,51 @@ def hex_to_signed(hex_str, bits):
如果输入的十六进制值超出指定位数的有符号整数范围，可能会返回错误的结果。
"""
unsigned_val = int(hex_str, 16)
mask = (1 << (bits - 1))
mask = 1 << (bits - 1)
if unsigned_val & mask:
return unsigned_val - (1 << bits)
else:
return unsigned_val
grayscale = transforms.Grayscale(num_output_channels=1)
resize = transforms.Resize((8, 8))
class ImageHashNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {"images": ("IMAGE", ),
"names": ("STRING",), },
"required": {
"images": ("IMAGE",),
"names": ("STRING",),
},
}
CATEGORY = "utils"
RETURN_TYPES = ('STRING',)
RETURN_TYPES = ("STRING",)
FUNCTION = "image_hash"
def image_hash(self, images,names):
def image_hash(self, images, names):
# 计算图像的哈希值
hashes = {}
from imagehash import ImageHash
file_names=[]
file_names = []
if isinstance(names, list):
file_names = names
else:
file_names = [names]
for i,image in enumerate(images):
# ComfyUI的IMAGE张量通常是[Batch, Height, Width, Channels]格式
# PyTorch transforms期望的是[Channels, Height, Width]格式
# 确保我们处理的是[Height, Width, Channels]格式
if len(image.shape) == 4: # 如果是[B, H, W, C]格式，取第一个
for i, image in enumerate(images):
# 确保处理的是 [H, W, C] 格式
if len(image.shape) == 4 and image.shape[0] == 1:
image = image.squeeze(0) # 移除批次维度，得到[H, W, C]
# 现在image应该是[H, W, C]格式
if len(image.shape) == 3 and image.shape[-1] in [1, 3]: # [Height, Width, Channels]
if len(image.shape) == 3 and image.shape[-1] in [
1,
3,
]: # [Height, Width, Channels]
tensor = image.permute(2, 0, 1) # 转换为 [Channels, Height, Width]
else:
tensor = image
@@ -242,13 +274,13 @@ class ImageHashNode:
pixels = binary_hash.squeeze(0).numpy()
mean = pixels.mean()
diff = pixels > mean
hash=hex_to_signed(str(ImageHash(diff)), 64)
hash = hex_to_signed(str(ImageHash(diff)), 64)
hashes.update({file_names[i]: hash})
#list to str
# list to str
return (json.dumps(hashes),)
class Resize2DivisibleImage:
class Resize2DivisibleImage:
@classmethod
def INPUT_TYPES(s):
return {
@@ -256,21 +288,28 @@ class Resize2DivisibleImage:
"image": ("IMAGE",),
"name": ("STRING",),
"target": ("INT", {"default": 32, "min": 2, "max": 2048, "step": 1}),
"max_size": ("INT", {"default": 1024, "min": 8, "max": 8192, "step": 1}), # 0 表示不限制
"max_size": (
"INT",
{"default": 1024, "min": 8, "max": 8192, "step": 1},
), # 0 表示不限制
},
}
CATEGORY = "utils"
RETURN_TYPES = ('IMAGE','STRING',)
RETURN_TYPES = (
"IMAGE",
"STRING",
)
FUNCTION = "divisible_resize_image"
def resize_image(self, image,target, max_size=0):
def resize_image(self, image, target, max_size=0):
"""
将图像的宽高调整到能被target整除的尺寸（保持宽高比例）
Args:
image: ComfyUI的IMAGE张量 [Batch, Height, Width, Channels]
target: 目标整除数
Returns:
调整后的图像
"""
@@ -300,26 +339,24 @@ class Resize2DivisibleImage:
image_permuted = image.permute(0, 3, 1, 2)
import torch.nn.functional as F
resized = F.interpolate(
image_permuted,
size=(new_h, new_w),
mode='bilinear',
align_corners=True
image_permuted, size=(new_h, new_w), mode="bilinear", align_corners=True
)
# 转换回 [B, H, W, C]
resized = resized.permute(0, 2, 3, 1).squeeze(0)
resized = resized.permute(0, 2, 3, 1)
return resized
def divisible_resize_image(self, image, name, target, max_size=0):
"""
将图像的宽高调整到能被target整除的尺寸
Args:
image: ComfyUI的IMAGE张量 [Batch, Height, Width, Channels]
name: 图像名称
target: 目标整除数
Returns:
调整后的图像和名称
"""
@@ -331,7 +368,8 @@ class Resize2DivisibleImage:
resized = []
for img in images:
resized.append(self.resize_image(img, target, max_size))
return (resized, name)
return (torch.cat(resized,dim=0), name)
class Base64ImageLoader:
@classmethod
@@ -339,12 +377,19 @@ class Base64ImageLoader:
return {
"required": {
"base64_string": ("STRING",),
"name": ("STRING",), },
"name": ("STRING",),
},
}
CATEGORY = "utils"
RETURN_TYPES = ("IMAGE", "MASK","STRING",)
RETURN_TYPES = (
"IMAGE",
"MASK",
"STRING",
)
FUNCTION = "load_base64_image"
def load_base64_image(self, base64_string,name):
def load_base64_image(self, base64_string, name):
"""
将base64编码的图像数据转换为PIL图像对象。
@@ -356,14 +401,131 @@ class Base64ImageLoader:
"""
import base64
from io import BytesIO
# 将base64编码的字符串解码为字节流
image_bytes = base64.b64decode(base64_string)
image = Image.open(BytesIO(image_bytes))
return process_image(image,name)
return process_image(image, name)
@classmethod
def IS_CHANGED(s,base64_string,name):
def IS_CHANGED(s, base64_string, name):
return float("NaN")
class GLMOCRNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"names": ("STRING",),
},
"optional": {
"prompt_type": (
["text_recognition", "formula_recognition", "table_recognition"],
),
},
}
CATEGORY = "utils"
RETURN_TYPES = ("STRING",)
FUNCTION = "ocr_extract"
# 缓存模型和处理器
_model = None
_processor = None
def ocr_extract(self, images, names, prompt_type="text_recognition"):
"""
使用GLM-OCR模型从图像中提取文字
Args:
images: ComfyUI的IMAGE张量 [Batch, Height, Width, Channels]
names: 图像名称
prompt_type: 提示类型 (text_recognition, formula_recognition, table_recognition)
Returns:
STRING: 提取的文字结果JSON字符串
"""
from transformers import AutoProcessor, AutoModelForImageTextToText
# 处理文件名
file_names = names if isinstance(names, list) else [names]
results = {}
# 提示词映射
prompt_map = {
"text_recognition": "Text Recognition:",
"formula_recognition": "Formula Recognition:",
"table_recognition": "Table Recognition:",
}
prompt = prompt_map.get(prompt_type, "Text Recognition:")
# 初始化或使用缓存的模型和处理器
if GLMOCRNode._model is None or GLMOCRNode._processor is None:
GLMOCRNode._processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
GLMOCRNode._model = AutoModelForImageTextToText.from_pretrained(
pretrained_model_name_or_path="zai-org/GLM-OCR",
torch_dtype="auto",
device_map="auto",
)
processor = GLMOCRNode._processor
model = GLMOCRNode._model
for i, image in enumerate(images):
# 确保处理的是 [H, W, C] 格式
if len(image.shape) == 4 and image.shape[0] == 1:
image = image.squeeze(0)
# 将ComfyUI的IMAGE张量转换为PIL图像
# image格式: [H, W, C], 值范围: [0, 1]
n = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(n, 0, 255).astype(np.uint8))
# 构建符合chat template格式的messages
messages = [
{
"role": "user",
"content": [
{"type": "image", "url": img},
{"type": "text", "text": prompt},
],
}
]
# 使用processor处理（支持PIL图像）
inputs = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
# 移除不需要的token_type_ids
inputs.pop("token_type_ids", None)
# 生成识别结果
generated_ids = model.generate(**inputs, max_new_tokens=8192)
# 解码输出（跳过输入ID部分）
output_text = processor.decode(
generated_ids[0][inputs["input_ids"].shape[1] :],
skip_special_tokens=True,
)
# 移除可能的特殊token标记（如 <|user|> 等）
import re
output_text = re.sub(r"<\|[^|]+\|>", "", output_text).strip()
# 使用文件名或序号作为键
key = file_names[i] if i < len(file_names) else f"image_{i}"
results[key] = output_text
# 返回JSON字符串
return (json.dumps(results, ensure_ascii=False),)
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
@@ -373,6 +535,7 @@ NODE_CLASS_MAPPINGS = {
"Image2Base64": Image2Base64,
"Resize2DivisibleImage": Resize2DivisibleImage,
"Base64ImageLoader": Base64ImageLoader,
"GLMOCRNode": GLMOCRNode,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -383,4 +546,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"Image2Base64": "Image to Base64",
"Base64ImageLoader": "Base64 Image Loader",
"Resize2DivisibleImage": "Resize to visible Image",
"GLMOCRNode": "GLM-OCR",
}