131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
import torch
|
|
import numpy as np
|
|
import cv2
|
|
|
|
try:
|
|
import easyocr
|
|
except Exception:
|
|
easyocr = None
|
|
|
|
_OCR_READER_CACHE = None
|
|
|
|
def get_ocr_reader(languages=["ch_sim", "en"]):
|
|
global _OCR_READER_CACHE
|
|
if _OCR_READER_CACHE is None:
|
|
if easyocr is None:
|
|
class DummyReader:
|
|
def readtext(self, img):
|
|
return []
|
|
_OCR_READER_CACHE = DummyReader()
|
|
else:
|
|
try:
|
|
_OCR_READER_CACHE = easyocr.Reader(languages, gpu=True)
|
|
except Exception:
|
|
_OCR_READER_CACHE = easyocr.Reader(languages, gpu=False)
|
|
return _OCR_READER_CACHE
|
|
|
|
|
|
class TextPreserveNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"blur_radius": ("INT", {
|
|
"default": 31,
|
|
"min": 1,
|
|
"max": 201,
|
|
"step": 2,
|
|
"display": "number",
|
|
}),
|
|
"noise_sigma": ("FLOAT", {
|
|
"default": 50.0,
|
|
"min": 0.0,
|
|
"max": 200.0,
|
|
"step": 1.0,
|
|
"display": "number",
|
|
}),
|
|
"mask_dilation": ("INT", {
|
|
"default": 10,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "number",
|
|
}),
|
|
"text_threshold": ("FLOAT", {
|
|
"default": 0.3,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.05,
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("processed_image", "text_mask")
|
|
FUNCTION = "process"
|
|
CATEGORY = "Custom/TextPreserve"
|
|
|
|
def process(self, image, blur_radius, noise_sigma, mask_dilation, text_threshold):
|
|
if blur_radius % 2 == 0:
|
|
blur_radius += 1
|
|
|
|
output_images = []
|
|
output_masks = []
|
|
|
|
reader = get_ocr_reader()
|
|
|
|
for i in range(image.shape[0]):
|
|
img_tensor = image[i]
|
|
img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
h, w, c = img_np.shape
|
|
if c == 4:
|
|
img_np = img_np[:, :, :3]
|
|
c = 3
|
|
|
|
blurred = cv2.GaussianBlur(img_np, (blur_radius, blur_radius), 0)
|
|
|
|
noise = np.random.normal(0, noise_sigma, (h, w, c))
|
|
noisy_bg = blurred.astype(np.float32) + noise
|
|
noisy_bg = np.clip(noisy_bg, 0, 255).astype(np.uint8)
|
|
|
|
results = reader.readtext(img_np)
|
|
|
|
mask = np.zeros((h, w), dtype=np.uint8)
|
|
|
|
for r in results:
|
|
bbox, _, prob = r
|
|
if prob < text_threshold:
|
|
continue
|
|
pts = np.array(bbox, dtype=np.int32)
|
|
cv2.fillPoly(mask, [pts], 255)
|
|
|
|
if mask_dilation > 0:
|
|
kernel = np.ones((mask_dilation, mask_dilation), np.uint8)
|
|
mask = cv2.dilate(mask, kernel, iterations=1)
|
|
|
|
mask_blurred = cv2.GaussianBlur(mask, (5, 5), 0)
|
|
|
|
mask_norm = mask_blurred.astype(np.float32) / 255.0
|
|
mask_norm = np.expand_dims(mask_norm, axis=2)
|
|
|
|
final_img = (
|
|
img_np.astype(np.float32) * mask_norm
|
|
+ noisy_bg.astype(np.float32) * (1.0 - mask_norm)
|
|
)
|
|
final_img = np.clip(final_img, 0, 255).astype(np.uint8)
|
|
|
|
out_tensor = torch.from_numpy(final_img.astype(np.float32) / 255.0)
|
|
output_images.append(out_tensor)
|
|
|
|
mask_tensor = torch.from_numpy(mask.astype(np.float32) / 255.0)
|
|
output_masks.append(mask_tensor)
|
|
|
|
final_images_batch = torch.stack(output_images)
|
|
final_masks_batch = torch.stack(output_masks)
|
|
|
|
return (final_images_batch, final_masks_batch) |