From aaefb465e290aa3cc83dc24aa8cbf14f2b33ad05 Mon Sep 17 00:00:00 2001 From: bai Date: Wed, 11 Feb 2026 17:14:01 +0800 Subject: [PATCH] =?UTF-8?q?```=20feat(nodes):=20=E6=B7=BB=E5=8A=A0GLM-OCR?= =?UTF-8?q?=E8=8A=82=E7=82=B9=E5=B9=B6=E4=BC=98=E5=8C=96=E5=9B=BE=E5=83=8F?= =?UTF-8?q?=E5=A4=84=E7=90=86=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增GLMOCRNode类，支持文本、公式和表格识别功能 - 优化process_image函数，修复字符串引号格式问题 - 改进CustomImageLoader类，支持批量图像处理和更好的URL处理 - 优化Image2Base64类，改进图像张量处理逻辑 - 添加VisionOutputEmbedding2JSON节点，支持视觉输出转JSON - 新增ImageHashNode节点，提供图像哈希计算功能 - 优化Resize2DivisibleImage节点，改进图像尺寸调整算法 - 新增Base64ImageLoader节点，支持base64图像加载 - 统一代码格式，修复数值精度和字符串引号一致性问题 ``` --- src/image_embeddings/nodes.py | 318 ++++++++++++++++++++++++++-------- 1 file changed, 241 insertions(+), 77 deletions(-) diff --git a/src/image_embeddings/nodes.py b/src/image_embeddings/nodes.py index 2fc5544..d13f27d 100644 --- a/src/image_embeddings/nodes.py +++ b/src/image_embeddings/nodes.py @@ -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", }