Spaces:
Build error
Build error
| 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) | |