File size: 3,204 Bytes
a6fb9f6 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import cv2
import numpy as np
MASK_VALUE = 255
ADAPTIVE_BLOCK_SIZE = 35
ADAPTIVE_OFFSET = 8
BACKGROUND_RING_WIDTH = 3
MIN_COMPONENT_AREA = 2
DEFAULT_MARGIN = 2
MAX_MARGIN = 10
POLARITY_AUTO = "Auto"
POLARITY_DARK = "Dark text"
POLARITY_LIGHT = "Light text"
POLARITY_CHOICES = [POLARITY_AUTO, POLARITY_DARK, POLARITY_LIGHT]
def character_mask(rgb_image, polygons, margin=DEFAULT_MARGIN, polarity=POLARITY_AUTO):
margin = int(margin)
if not 0 <= margin <= MAX_MARGIN:
raise ValueError(f"Margin must be between 0 and {MAX_MARGIN}")
if polarity not in POLARITY_CHOICES:
raise ValueError("Invalid text polarity")
gray = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY)
height, width = gray.shape
mask = np.zeros_like(gray)
padding = max(ADAPTIVE_BLOCK_SIZE // 2, BACKGROUND_RING_WIDTH)
ring_size = BACKGROUND_RING_WIDTH * 2 + 1
ring_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ring_size, ring_size))
for polygon in polygons:
points = np.asarray(polygon, dtype=np.float32).reshape(-1, 2)
if len(points) < 3 or not np.isfinite(points).all():
continue
points = np.rint(points).astype(np.int32)
x, y, box_width, box_height = cv2.boundingRect(points)
left, top = max(0, x - padding), max(0, y - padding)
right = min(width, x + box_width + padding)
bottom = min(height, y + box_height + padding)
if left >= right or top >= bottom:
continue
crop = gray[top:bottom, left:right]
region = np.zeros_like(crop)
cv2.fillPoly(region, [points - (left, top)], MASK_VALUE)
inside = region != 0
if not np.any(inside):
continue
selected_polarity = polarity
if selected_polarity == POLARITY_AUTO:
ring = (cv2.dilate(region, ring_kernel) != 0) & ~inside
background = np.median(crop[ring] if np.any(ring) else crop[inside])
split, _ = cv2.threshold(
crop[inside], 0, MASK_VALUE, cv2.THRESH_BINARY | cv2.THRESH_OTSU
)
selected_polarity = POLARITY_DARK if background > split else POLARITY_LIGHT
dark_text = selected_polarity == POLARITY_DARK
threshold_type = cv2.THRESH_BINARY_INV if dark_text else cv2.THRESH_BINARY
offset = ADAPTIVE_OFFSET if dark_text else -ADAPTIVE_OFFSET
foreground = cv2.adaptiveThreshold(
crop,
MASK_VALUE,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
threshold_type,
ADAPTIVE_BLOCK_SIZE,
offset,
)
foreground = cv2.bitwise_and(foreground, region)
count, labels, stats, _ = cv2.connectedComponentsWithStats(foreground, connectivity=8)
keep = np.zeros(count, dtype=np.uint8)
keep[1:] = (stats[1:, cv2.CC_STAT_AREA] >= MIN_COMPONENT_AREA) * MASK_VALUE
cleaned = keep[labels]
target = mask[top:bottom, left:right]
np.maximum(target, cleaned, out=target)
if margin:
kernel_size = margin * 2 + 1
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
mask = cv2.dilate(mask, kernel)
return mask
|