File size: 2,673 Bytes
9dce458 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
from typing import List
import cv2
import numpy as np
from .text_mask_utils import complete_mask_fill, complete_mask
from ..utils import TextBlock, Quadrilateral
from ..utils.bubble import is_ignore
async def dispatch(text_regions: List[TextBlock], raw_image: np.ndarray, raw_mask: np.ndarray, method: str = 'fit_text', dilation_offset: int = 0, ignore_bubble: int = 0, verbose: bool = False,kernel_size:int=3) -> np.ndarray:
# Larger sized mask images will probably have crisper and thinner mask segments due to being able to fit the text pixels better
# so we dont want to size them down as much to not lose information
scale_factor = max(min((raw_mask.shape[0] - raw_image.shape[0] / 3) / raw_mask.shape[0], 1), 0.5)
img_resized = cv2.resize(raw_image, (int(raw_image.shape[1] * scale_factor), int(raw_image.shape[0] * scale_factor)), interpolation = cv2.INTER_LINEAR)
mask_resized = cv2.resize(raw_mask, (int(raw_image.shape[1] * scale_factor), int(raw_image.shape[0] * scale_factor)), interpolation = cv2.INTER_LINEAR)
mask_resized[mask_resized > 0] = 255
textlines = []
for region in text_regions:
for l in region.lines:
q = Quadrilateral(l * scale_factor, '', 0)
textlines.append(q)
final_mask = complete_mask(img_resized, mask_resized, textlines, dilation_offset=dilation_offset,kernel_size=kernel_size) if method == 'fit_text' else complete_mask_fill([txtln.aabb.xywh for txtln in textlines])
if final_mask is None:
final_mask = np.zeros((raw_image.shape[0], raw_image.shape[1]), dtype = np.uint8)
else:
final_mask = cv2.resize(final_mask, (raw_image.shape[1], raw_image.shape[0]), interpolation = cv2.INTER_LINEAR)
final_mask[final_mask > 0] = 255
if ignore_bubble < 1 or ignore_bubble > 50:
return final_mask
# bubble
kernel_size = int(max(final_mask.shape) * 0.025) # 选择一个合适的核大小
kernel = np.ones((kernel_size, kernel_size), np.uint8)
final_mask = cv2.dilate(final_mask, kernel, iterations=1) # 根据需要调整迭代次数
# border
contours, _ = cv2.findContours(final_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for cnt in contours:
temp_mask = np.zeros_like(final_mask)
# rect min
x, y, w, h = cv2.boundingRect(cnt)
cv2.rectangle(temp_mask, (x, y), (x + w, y + h), 255, -1)
# get textblock
textblock=cv2.bitwise_and(raw_image, raw_image, mask=temp_mask)
if is_ignore(textblock, ignore_bubble):
cv2.drawContours(final_mask, [cnt], -1, 0, -1)
return final_mask
|