diff --git a/nodes/anytext.py b/nodes/anytext.py deleted file mode 100644 index e950b16..0000000 --- a/nodes/anytext.py +++ /dev/null @@ -1,60 +0,0 @@ -from modelscope.pipelines import pipeline -import json -import os -import re -import base64 -from io import BytesIO -import requests -from retry import retry -import boto3 -from PIL import Image -import numpy as np -import torch -import cv2 - -class AnyTextGen: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ori_image": ("IMAGE",), - "mask_image":("IMAGE",), - "prompt":("STRING",) - } - } - - RETURN_TYPES = ("IMAGE") - RETURN_NAMES = ("Converted Image") - FUNCTION = "forward" - CATEGORY = "aws" - OUTPUT_NODE = True - - def any_text_edit_image(self,image_input): - pipe = pipeline('my-anytext-task', model='damo/cv_anytext_text_generation_editing', model_revision='v1.1.1', use_fp16=True, use_translator=False) - img_save_folder = "SaveImages" - params = { - "show_debug": True, - "image_count": 1, - "ddim_steps": 20, - } - - mode = 'text-editing' - input_data = { - "prompt": prompt, - "seed": 8943410, - "draw_pos": ori_image, - "ori_image": mask_image - } - results, rtn_code, rtn_warning, debug_info = pipe(input_data, mode=mode, **params) - image = torch.from_numpy(np.array(results[0]).astype(np.float32) / 255.0).unsqueeze(0) - return image - - @retry(tries=3) - def forward(self, image): - return self.any_text_edit_image(image) - - - -NODE_CLASS_MAPPINGS = { - "AnyText Gen": AnyTextGen -} \ No newline at end of file diff --git a/nodes/textract.py b/nodes/textract.py index 4081f83..2b7bf20 100644 --- a/nodes/textract.py +++ b/nodes/textract.py @@ -101,7 +101,7 @@ class ImageOCRByTextract: return self.ocr_by_textract(image) - +### for anyText Nodes class ImageOCRByTextractV2: @classmethod @@ -198,8 +198,121 @@ class ImageOCRByTextractV2: return self.ocr_by_textract(image) +## for layer style nodes +class ImageOCRByTextractV3: + + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + return {"required": + { + "image": (sorted(files), {"image_upload": True}), + }, + } + + RETURN_TYPES = ("STRING","STRING","STRING","STRING","STRING","IMAGE") + RETURN_NAMES = ("Texts","x_offsets","y_offsets","widths","heights","AnyText Mask Image") + FUNCTION = "forward" + CATEGORY = "aws" + OUTPUT_NODE = True + + def ocr_by_textract(self,image_input): + ori_image_path = folder_paths.get_annotated_filepath(image_input) + + ## 转换ori_image为tensor张量 + pil_image = Image.open(ori_image_path) + transform = transforms.Compose([ + transforms.ToTensor(), # 将 PIL Image 转换为 tensor，并将像素值归一化到 [0, 1] + ]) + image_input = transform(pil_image) + + ## image input已经是标准comfyui的image张量格式 + numpy_image = (image_input[0] * 255.0).clamp(0, 255).numpy() + image = image_input[0] * 255.0 + image = Image.fromarray(image.clamp(0, 255).numpy().round().astype(np.uint8)) + # 获取图像原始尺寸 + img_width, img_height = image.size + + # 调用Textract DetectDocumentText函数 + byte_stream = io.BytesIO() + image.save(byte_stream, format='PNG') + byte_image = byte_stream.getvalue() + response = textract.detect_document_text(Document={'Bytes': byte_image}) + + # 初始化结果列表 + result = [] + # 创建一个与原始图像大小相同的黑色遮罩图像 + #mask = np.zeros_like(numpy_image) + mask = np.ones_like(numpy_image) * 255 + all_text="" + x_offsets=[] + y_offsets=[] + widths=[] + heights=[] + + + + # 提取文本和边界框信息 + for item in response['Blocks']: + if item['BlockType'] == 'LINE': + text = item['Text'] + box = item['Geometry']['BoundingBox'] + + left = int(box['Left'] * img_width) + x_offsets.append(str(left)) + + top = int(box['Top'] * img_height) + y_offsets.append(str(top)) + + width = int(box['Width'] * img_width) + widths.append(str(width)) + + height = int(box['Height']* img_height) + heights.append(str(height)) + + # 将信息添加到结果列表中 + result.append(text) + # 对每个文本信息框绘制mask遮罩 + # 指定矩形框的左上角和右下角坐标 + temp_mask = np.zeros_like(numpy_image) + x1, y1 = int(left), int(top) + x2, y2 = int(left + width), int(top + height) + # 在遮罩图像上绘制黑色矩形框 + cv2.rectangle(temp_mask, (x1, y1), (x2, y2), (255, 255, 255), -1) + + + # 将遮罩图像和原始图像进行位运算,生成遮罩后的图像 + masked_img = np.where(mask == 0, 0, numpy_image) + #masked_img = cv2.bitwise_and(numpy_image, mask) + masked_img = torch.from_numpy(np.array(masked_img).astype(np.float32) / 255.0).unsqueeze(0) + #masked_img.save(temp_img_path) + # 将 PyTorch 张量转换为 PIL 图像 + to_pil = transforms.ToPILImage() + pil_image = to_pil(masked_img.squeeze(0)) + # 保存 PIL 图像 + pil_image.save(temp_img_path) + + ###汇总结果输出 + all_text="|".join(result) + x_offsets="|".join(x_offsets) + y_offsets="|".join(y_offsets) + widths="|".join(widths) + heights="|".join(heights) + + print("result",result) + + return all_text ,x_offsets,y_offsets,widths,heights, masked_img + + + + @retry(tries=MAX_RETRY) + def forward(self, image): + return self.ocr_by_textract(image) + NODE_CLASS_MAPPINGS = { "Image OCR By Textract": ImageOCRByTextract, - "Image OCR By Textract V2":ImageOCRByTextractV2 + "Image OCR By Textract V2":ImageOCRByTextractV2, + "Image OCR By Textract V3":ImageOCRByTextractV3 }