eKeerar
108 lines
3.0 KiB
Python
108 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
|
|
def connected_components(mask: np.ndarray) -> np.ndarray:
|
|
h, w = mask.shape
|
|
labels = np.zeros((h, w), dtype=np.int32)
|
|
current = 0
|
|
|
|
for y in range(h):
|
|
for x in range(w):
|
|
if not mask[y, x] or labels[y, x] != 0:
|
|
continue
|
|
current += 1
|
|
queue: deque[tuple[int, int]] = deque([(x, y)])
|
|
labels[y, x] = current
|
|
while queue:
|
|
cx, cy = queue.popleft()
|
|
for nx, ny in (
|
|
(cx + 1, cy),
|
|
(cx - 1, cy),
|
|
(cx, cy + 1),
|
|
(cx, cy - 1),
|
|
):
|
|
if nx < 0 or ny < 0 or nx >= w or ny >= h:
|
|
continue
|
|
if mask[ny, nx] and labels[ny, nx] == 0:
|
|
labels[ny, nx] = current
|
|
queue.append((nx, ny))
|
|
return labels
|
|
|
|
|
|
def main() -> None:
|
|
if len(sys.argv) != 3:
|
|
raise SystemExit("usage: make_hairmask.py <input-head.png> <output-mask.png>")
|
|
|
|
src = Path(sys.argv[1])
|
|
out = Path(sys.argv[2])
|
|
|
|
image = Image.open(src).convert("RGBA")
|
|
arr = np.asarray(image)
|
|
rgb = arr[:, :, :3].astype(np.int16)
|
|
alpha = arr[:, :, 3]
|
|
r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
|
|
|
|
opaque = alpha > 32
|
|
brown_hair = (
|
|
opaque
|
|
& (r > 34)
|
|
& (g > 18)
|
|
& (r >= g - 8)
|
|
& (g >= b - 4)
|
|
& ((r - b) > 22)
|
|
& ((r < 205) | (g < 138) | (b < 105))
|
|
)
|
|
dark_hair_lines = opaque & (r > 18) & (r < 115) & (g < 100) & (b < 85) & ((r - b) > 8)
|
|
mask = brown_hair | dark_hair_lines
|
|
|
|
labels = connected_components(mask)
|
|
keep = np.zeros_like(mask)
|
|
for label in range(1, labels.max() + 1):
|
|
ys, xs = np.where(labels == label)
|
|
if xs.size == 0:
|
|
continue
|
|
area = xs.size
|
|
width = int(xs.max() - xs.min() + 1)
|
|
height = int(ys.max() - ys.min() + 1)
|
|
left = int(xs.min())
|
|
right = int(xs.max())
|
|
top = int(ys.min())
|
|
bottom = int(ys.max())
|
|
h, w = arr.shape[:2]
|
|
touches_hair_zone = (
|
|
left < w * 0.25
|
|
or right > w * 0.75
|
|
or top < h * 0.32
|
|
or bottom > h * 0.78
|
|
)
|
|
if area >= 5000 or (
|
|
area >= 1200
|
|
and touches_hair_zone
|
|
and (width >= 45 or height >= 45)
|
|
):
|
|
keep[ys, xs] = True
|
|
|
|
mask_img = Image.fromarray((keep.astype(np.uint8) * 255), "L")
|
|
mask_img = mask_img.filter(ImageFilter.MaxFilter(5))
|
|
mask_arr = np.asarray(mask_img) > 0
|
|
mask_arr &= opaque
|
|
|
|
out_arr = np.zeros_like(arr)
|
|
out_arr[:, :, :3] = 255
|
|
out_arr[:, :, 3] = (mask_arr.astype(np.uint8) * 255)
|
|
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
Image.fromarray(out_arr, "RGBA").save(out)
|
|
print(f"Wrote {out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|