Spaces:
Running on Zero
Running on Zero
| # coding=utf-8 | |
| """Turn a gr.ImageEditor value into a background image, a stroke mask and a pixel box.""" | |
| from __future__ import annotations | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| # Diff threshold for the fallback path; the frontend re-encodes images as webp, | |
| # so compression noise has to be filtered out. | |
| DIFF_THRESHOLD = 32 | |
| MIN_COMPONENT_AREA = 24 | |
| def keep_large_components(mask, min_area=MIN_COMPONENT_AREA): | |
| """Drop small blobs caused by compression noise.""" | |
| m = np.asarray(mask).astype(np.uint8) | |
| n, labels, stats, _ = cv2.connectedComponentsWithStats(m, connectivity=8) | |
| out = np.zeros(m.shape, dtype=bool) | |
| for i in range(1, n): | |
| if stats[i, cv2.CC_STAT_AREA] >= min_area: | |
| out |= labels == i | |
| return out | |
| def _layer_alpha_mask(layer, size): | |
| if layer is None: | |
| return None | |
| if layer.size != size: | |
| layer = layer.resize(size, Image.BILINEAR) | |
| return np.asarray(layer.convert("RGBA"))[..., 3] > 10 | |
| def editor_mask(editor_value): | |
| """Return (background RGB, stroke mask). | |
| Prefer the brush layer alpha; fall back to a background/composite diff when the | |
| frontend does not send the layer. | |
| """ | |
| if not editor_value: | |
| return None, None | |
| if not isinstance(editor_value, dict): | |
| bg = editor_value.convert("RGB") | |
| return bg, np.zeros((bg.size[1], bg.size[0]), dtype=bool) | |
| bg = editor_value.get("background") | |
| if bg is None: | |
| return None, None | |
| bg = bg.convert("RGB") | |
| h, w = bg.size[1], bg.size[0] | |
| layer_mask = np.zeros((h, w), dtype=bool) | |
| for layer in (editor_value.get("layers") or []): | |
| m = _layer_alpha_mask(layer, bg.size) | |
| if m is not None: | |
| layer_mask |= m | |
| if layer_mask.any(): | |
| return bg, keep_large_components(layer_mask) | |
| comp = editor_value.get("composite") | |
| if comp is None: | |
| return bg, np.zeros((h, w), dtype=bool) | |
| comp = comp.convert("RGB") | |
| if comp.size != bg.size: | |
| comp = comp.resize(bg.size, Image.BILINEAR) | |
| diff = np.abs(np.asarray(bg, np.int16) - np.asarray(comp, np.int16)).max(axis=2) | |
| return bg, keep_large_components(diff > DIFF_THRESHOLD) | |
| def mask_to_box(mask, pad_px=0): | |
| """Bounding box of the mask as (x0, y0, x1, y1), right/bottom exclusive.""" | |
| ys, xs = np.nonzero(np.asarray(mask)) | |
| if len(xs) == 0: | |
| return None | |
| pad = int(max(0, pad_px)) | |
| return (int(xs.min()) - pad, int(ys.min()) - pad, | |
| int(xs.max()) + 1 + pad, int(ys.max()) + 1 + pad) | |
| def clamp_box(box, img_w, img_h): | |
| """Clamp a box to the image bounds; None if it is degenerate.""" | |
| if box is None: | |
| return None | |
| x0, y0, x1, y1 = [float(v) for v in box] | |
| if x1 < x0: | |
| x0, x1 = x1, x0 | |
| if y1 < y0: | |
| y0, y1 = y1, y0 | |
| x0, x1 = max(0.0, min(x0, img_w)), max(0.0, min(x1, img_w)) | |
| y0, y1 = max(0.0, min(y0, img_h)), max(0.0, min(y1, img_h)) | |
| if x1 - x0 < 2 or y1 - y0 < 2: | |
| return None | |
| return (int(round(x0)), int(round(y0)), int(round(x1)), int(round(y1))) | |
| def manual_box(values, img_w, img_h): | |
| """Parse manual x0/y0/x1/y1 inputs; None when invalid.""" | |
| try: | |
| x0, y0, x1, y1 = [float(v) for v in values] | |
| except (TypeError, ValueError): | |
| return None | |
| if x1 <= x0 or y1 <= y0: | |
| return None | |
| return clamp_box((x0, y0, x1, y1), img_w, img_h) | |