Spaces:
Build error
Build error
File size: 3,856 Bytes
3c2cd23 | 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 86 87 88 89 90 91 92 93 94 95 | from __future__ import annotations
import cv2
import numpy as np
from .models import DocumentSpec, MaskStroke, Point, TextRegion
from .shape_recognition import shape_source_mask
def build_mask(rgba: np.ndarray, document: DocumentSpec) -> np.ndarray:
height, width = rgba.shape[:2]
combined = np.zeros((height, width), dtype=np.uint8)
for region in document.regions:
if region.remove_text:
combined = cv2.bitwise_or(combined, region_mask(rgba, region))
for shape in document.shapes:
if shape.remove_source:
shape_mask = shape_source_mask(shape, width, height)
padding = max(0, int(round(shape.mask_padding)))
if padding:
kernel = cv2.getStructuringElement(
cv2.MORPH_ELLIPSE,
(padding * 2 + 1, padding * 2 + 1),
)
shape_mask = cv2.dilate(shape_mask, kernel)
combined = cv2.bitwise_or(combined, shape_mask)
for stroke in document.mask_strokes:
apply_stroke(combined, stroke)
return combined
def region_mask(rgba: np.ndarray, region: TextRegion) -> np.ndarray:
height, width = rgba.shape[:2]
polygon = np.array([[point.x, point.y] for point in region.source_quad], dtype=np.int32)
bounds = cv2.boundingRect(polygon)
x, y, w, h = bounds
x0, y0 = max(0, x), max(0, y)
x1, y1 = min(width, x + w), min(height, y + h)
output = np.zeros((height, width), dtype=np.uint8)
if x1 <= x0 or y1 <= y0:
return output
crop = rgba[y0:y1, x0:x1, :3]
lab = cv2.cvtColor(crop, cv2.COLOR_RGB2LAB).astype(np.float32)
local_polygon = polygon - np.array([x0, y0])
polygon_mask = np.zeros(crop.shape[:2], dtype=np.uint8)
cv2.fillPoly(polygon_mask, [local_polygon], 255)
perimeter = cv2.morphologyEx(polygon_mask, cv2.MORPH_GRADIENT, np.ones((3, 3), np.uint8))
border_pixels = lab[perimeter > 0]
background = np.median(border_pixels, axis=0) if border_pixels.size else np.median(lab.reshape(-1, 3), axis=0)
delta = np.linalg.norm(lab - background, axis=2)
inside = delta[polygon_mask > 0]
if inside.size:
otsu_threshold, _ = cv2.threshold(
np.clip(inside, 0, 255).astype(np.uint8),
0,
255,
cv2.THRESH_BINARY + cv2.THRESH_OTSU,
)
threshold = max(8.0, float(otsu_threshold) * 0.72)
else:
threshold = 8.0
foreground = ((delta >= threshold) & (polygon_mask > 0)).astype(np.uint8) * 255
count, labels, stats, _ = cv2.connectedComponentsWithStats(foreground, 8)
filtered = np.zeros_like(foreground)
max_area = max(9, int(w * h * 0.22))
min_area = 1
for index in range(1, count):
area = stats[index, cv2.CC_STAT_AREA]
component_height = stats[index, cv2.CC_STAT_HEIGHT]
if min_area <= area <= max_area and component_height >= 1:
filtered[labels == index] = 255
if cv2.countNonZero(filtered) < max(3, w * h * 0.005):
filtered = foreground
padding = max(1, int(round(region.mask_padding)))
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (padding * 2 + 1, padding * 2 + 1))
filtered = cv2.dilate(filtered, kernel)
output[y0:y1, x0:x1] = filtered
return output
def apply_stroke(mask: np.ndarray, stroke: MaskStroke) -> None:
value = 255 if stroke.operation == "add" else 0
radius = max(1, int(round(stroke.radius)))
points = [(int(round(point.x)), int(round(point.y))) for point in stroke.points]
if len(points) == 1:
cv2.circle(mask, points[0], radius, value, -1, lineType=cv2.LINE_AA)
return
for start, end in zip(points, points[1:]):
cv2.line(mask, start, end, value, radius * 2, lineType=cv2.LINE_AA)
cv2.circle(mask, end, radius, value, -1, lineType=cv2.LINE_AA)
|