261 lines
7.2 KiB
Python
261 lines
7.2 KiB
Python
import folder_paths
|
|
import json
|
|
import cv2
|
|
import easyocr
|
|
import os
|
|
import logging
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
logger = logging.getLogger("ComfyUI-EasyOCR")
|
|
model_dir_name = "EasyOCR"
|
|
|
|
lang_list = {
|
|
"English": "en",
|
|
"简体中文": "ch_sim",
|
|
"繁體中文": "ch_tra",
|
|
"العربية": "ar",
|
|
"Azərbaycan": "az",
|
|
"Euskal": "eu",
|
|
"Bosanski": "bs",
|
|
"Български": "bg",
|
|
"Català": "ca",
|
|
"Hrvatski": "hr",
|
|
"Čeština": "cs",
|
|
"Dansk": "da",
|
|
"Nederlands": "nl",
|
|
"Eesti": "et",
|
|
"Suomi": "fi",
|
|
"Français": "fr",
|
|
"Galego": "gl",
|
|
"Deutsch": "de",
|
|
"Ελληνικά": "el",
|
|
"עברית": "he",
|
|
"हिन्दी": "hi",
|
|
"Magyar": "hu",
|
|
"Íslenska": "is",
|
|
"Indonesia": "id",
|
|
"Italiano": "it",
|
|
"日本語": "ja",
|
|
"한국어": "ko",
|
|
"Latviešu": "lv",
|
|
"Lietuvių": "lt",
|
|
"Македонски": "mk",
|
|
"Norsk": "no",
|
|
"Polski": "pl",
|
|
"Português": "pt",
|
|
"Română": "ro",
|
|
"Русский": "ru",
|
|
"Српски": "sr",
|
|
"Slovenčina": "sk",
|
|
"Slovenščina": "sl",
|
|
"Español": "es",
|
|
"Svenska": "sv",
|
|
"ไทย": "th",
|
|
"Türkçe": "tr",
|
|
"Українська": "uk",
|
|
"Tiếng Việt": "vi",
|
|
}
|
|
|
|
def get_lang_list():
|
|
result = []
|
|
for key, value in lang_list.items():
|
|
result.append(key)
|
|
return result
|
|
|
|
|
|
def get_classes(label):
|
|
label = label.lower()
|
|
labels = label.split(",")
|
|
result = []
|
|
for l in labels:
|
|
for key, value in lang_list.items():
|
|
if l == value:
|
|
result.append(value)
|
|
break
|
|
return result
|
|
|
|
|
|
def get_classes2(label):
|
|
label = label.lower()
|
|
labels = label.split(",")
|
|
result = []
|
|
for l in labels:
|
|
for key, value in lang_list.items():
|
|
if l == key:
|
|
result.append(value)
|
|
break
|
|
return result
|
|
|
|
|
|
def plot_boxes_to_image(image_pil, tgt):
|
|
H, W = tgt["size"]
|
|
result = tgt["result"]
|
|
|
|
res_mask = []
|
|
res_image = []
|
|
|
|
box_color = (255, 0, 0) # Red color for the box
|
|
text_color = (255, 255, 255) # White color for the text
|
|
|
|
draw = ImageDraw.Draw(image_pil)
|
|
|
|
# Get the current file path and use it to create a relative path to the font file
|
|
current_file_path = os.path.dirname(os.path.abspath(__file__))
|
|
font_path = os.path.join(current_file_path, "docs", "PingFang Regular.ttf")
|
|
font_size = 20
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
|
|
labelme_data = {
|
|
"version": "4.5.6",
|
|
"flags": {},
|
|
"shapes": [],
|
|
"imagePath": None,
|
|
"imageData": None,
|
|
"imageHeight": H,
|
|
"imageWidth": W,
|
|
}
|
|
|
|
for item in result:
|
|
formatted_points, label, threshold = item
|
|
|
|
x1, y1 = formatted_points[0]
|
|
x2, y2 = formatted_points[2]
|
|
threshold = round(threshold, 2)
|
|
|
|
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
|
points = [[x1, y1], [x2, y2]]
|
|
|
|
# Save labelme json
|
|
shape = {
|
|
"label": label,
|
|
"points": points,
|
|
"group_id": None,
|
|
"shape_type": "rectangle",
|
|
"flags": {},
|
|
}
|
|
labelme_data["shapes"].append(shape)
|
|
|
|
# Change label
|
|
label = label + ":" + str(threshold)
|
|
shape["threshold"] = str(threshold)
|
|
|
|
# Draw rectangle on the image using PIL
|
|
draw.rectangle([(x1, y1), (x2, y2)], outline=box_color, width=3)
|
|
|
|
# Draw label on the image using PIL
|
|
text_bbox = draw.textbbox((x1, y1), label, font=font)
|
|
text_width = text_bbox[2] - text_bbox[0]
|
|
text_height = text_bbox[3] - text_bbox[1]
|
|
|
|
label_ymin = max(y1, text_height + 10)
|
|
draw.rectangle(
|
|
[(x1, y1 - text_height - 10), (x1 + text_width, y1)], fill=box_color
|
|
)
|
|
draw.text((x1, y1 - text_height - 10), label, font=font, fill=text_color)
|
|
|
|
# Draw mask
|
|
mask = np.zeros((H, W, 1), dtype=np.uint8)
|
|
cv2.rectangle(mask, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 255), -1)
|
|
mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
|
|
res_mask.append(mask_tensor)
|
|
|
|
if len(res_mask) == 0:
|
|
mask = np.zeros((H, W, 1), dtype=np.uint8)
|
|
mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
|
|
res_mask.append(mask_tensor)
|
|
|
|
# Convert the PIL image back to a numpy array
|
|
image_with_boxes = np.array(image_pil)
|
|
|
|
# Convert the modified image to a torch tensor
|
|
image_with_boxes_tensor = torch.from_numpy(
|
|
image_with_boxes.astype(np.float32) / 255.0
|
|
)
|
|
image_with_boxes_tensor = torch.unsqueeze(image_with_boxes_tensor, 0)
|
|
res_image.append(image_with_boxes_tensor)
|
|
|
|
return res_image, res_mask, labelme_data
|
|
|
|
|
|
class ApplyEasyOCR:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"gpu": (
|
|
"BOOLEAN",
|
|
{"default": True},
|
|
),
|
|
"detect": (
|
|
["choose", "input"],
|
|
{"default": "choose"},
|
|
),
|
|
"language_list": (
|
|
get_lang_list(),
|
|
{"default": "English"},
|
|
),
|
|
"language_name": (
|
|
"STRING",
|
|
{"default": "ch_sim,en", "multiline": False},
|
|
),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "ComfyUI-EasyOCR"
|
|
FUNCTION = "main"
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"MASK",
|
|
"JSON",
|
|
)
|
|
|
|
def main(self, image, gpu, detect, language_list, language_name):
|
|
res_images = []
|
|
res_masks = []
|
|
res_labels = []
|
|
|
|
for item in image:
|
|
image_pil = Image.fromarray(np.clip(255.0 * item.cpu().numpy(), 0, 255).astype(np.uint8)).convert("RGB")
|
|
|
|
language = None
|
|
if detect == "choose":
|
|
language = get_classes2(language_list)
|
|
else:
|
|
language = get_classes(language_name)
|
|
|
|
model_storage_directory = os.path.join(folder_paths.models_dir, model_dir_name)
|
|
if not os.path.exists(model_storage_directory):
|
|
os.makedirs(model_storage_directory)
|
|
|
|
reader = easyocr.Reader(language, model_storage_directory=model_storage_directory,gpu=gpu)
|
|
result = reader.readtext(np.array(image_pil))
|
|
|
|
size = image_pil.size
|
|
pred_dict = {
|
|
"size": [size[1], size[0]],
|
|
"result":result
|
|
}
|
|
|
|
image_tensor, mask_tensor, labelme_data = plot_boxes_to_image(image_pil, pred_dict)
|
|
|
|
res_images.extend(image_tensor)
|
|
res_masks.extend(mask_tensor)
|
|
res_labels.append(labelme_data)
|
|
|
|
if len(res_images) == 0:
|
|
res_images.extend(item)
|
|
if len(res_masks) == 0:
|
|
mask = np.zeros((height, width, 1), dtype=np.uint8)
|
|
empty_mask = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
|
|
res_masks.extend(empty_mask)
|
|
|
|
return (
|
|
torch.cat(res_images, dim=0),
|
|
torch.cat(res_masks, dim=0),
|
|
res_labels,
|
|
)
|