| 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 |
|
|
