Spaces:
Running on Zero
Running on Zero
File size: 3,380 Bytes
3117cf2 | 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 96 97 98 99 100 101 102 103 104 | # 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)
|