116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
def parse_background_color(value: str) -> tuple[int, int, int]:
|
|
text = value.strip()
|
|
if text.startswith("#"):
|
|
hex_text = text[1:]
|
|
if len(hex_text) != 6:
|
|
raise ValueError("background_color 使用 #RRGGBB 时必须是 6 位十六进制")
|
|
try:
|
|
return tuple(int(hex_text[i : i + 2], 16) for i in (0, 2, 4)) # type: ignore[return-value]
|
|
except ValueError as e:
|
|
raise ValueError(f"background_color 不是有效十六进制颜色: {value}") from e
|
|
|
|
parts = [p.strip() for p in text.split(",")]
|
|
if len(parts) != 3:
|
|
raise ValueError("background_color 必须为 #RRGGBB 或 R,G,B")
|
|
try:
|
|
rgb = tuple(int(p) for p in parts)
|
|
except ValueError as e:
|
|
raise ValueError(f"background_color 不是有效 RGB 颜色: {value}") from e
|
|
if any(c < 0 or c > 255 for c in rgb):
|
|
raise ValueError("background_color 的 RGB 分量必须在 0~255")
|
|
return rgb # type: ignore[return-value]
|
|
|
|
|
|
def align_size_to_input(width: int, height: int, max_side: int) -> tuple[int, int]:
|
|
run_w, run_h = width, height
|
|
if max(width, height) > max_side:
|
|
scale = max_side / float(max(width, height))
|
|
run_w = int(round(width * scale))
|
|
run_h = int(round(height * scale))
|
|
|
|
run_w = max(8, run_w - (run_w % 8))
|
|
run_h = max(8, run_h - (run_h % 8))
|
|
return run_w, run_h
|
|
|
|
|
|
def composite_rgba_on_background(
|
|
image: Image.Image,
|
|
background_rgb: tuple[int, int, int],
|
|
size: tuple[int, int],
|
|
) -> Image.Image:
|
|
rgba = image.convert("RGBA")
|
|
if rgba.size != size:
|
|
rgba = rgba.resize(size, resample=Image.BICUBIC)
|
|
|
|
bg = Image.new("RGB", size, background_rgb)
|
|
alpha = rgba.getchannel("A")
|
|
bg.paste(rgba.convert("RGB"), mask=alpha)
|
|
return bg
|
|
|
|
|
|
def remove_background_by_color(
|
|
image: Image.Image,
|
|
background_rgb: tuple[int, int, int],
|
|
tolerance: int,
|
|
output_size: tuple[int, int] | None = None,
|
|
) -> Image.Image:
|
|
rgba = image.convert("RGBA")
|
|
if output_size is not None and rgba.size != output_size:
|
|
rgba = rgba.resize(output_size, resample=Image.BICUBIC)
|
|
|
|
arr = np.asarray(rgba, dtype=np.uint8).copy()
|
|
rgb = arr[:, :, :3].astype(np.int16)
|
|
bg = np.array(background_rgb, dtype=np.int16).reshape(1, 1, 3)
|
|
diff = np.abs(rgb - bg)
|
|
bg_like = np.max(diff, axis=2) <= int(tolerance)
|
|
bg_mask = _edge_connected_mask(bg_like)
|
|
arr[bg_mask, 3] = 0
|
|
return Image.fromarray(arr, mode="RGBA")
|
|
|
|
|
|
def list_png_frames(frame_dir: Path) -> list[Path]:
|
|
frames = [p for p in frame_dir.iterdir() if p.is_file() and p.suffix.lower() == ".png"]
|
|
return sorted(frames, key=lambda p: p.name)
|
|
|
|
|
|
def _edge_connected_mask(mask: np.ndarray) -> np.ndarray:
|
|
h, w = mask.shape
|
|
visited = np.zeros((h, w), dtype=bool)
|
|
queue: deque[tuple[int, int]] = deque()
|
|
|
|
for x in range(w):
|
|
if mask[0, x]:
|
|
queue.append((0, x))
|
|
if mask[h - 1, x]:
|
|
queue.append((h - 1, x))
|
|
for y in range(h):
|
|
if mask[y, 0]:
|
|
queue.append((y, 0))
|
|
if mask[y, w - 1]:
|
|
queue.append((y, w - 1))
|
|
|
|
while queue:
|
|
y, x = queue.popleft()
|
|
if visited[y, x] or not mask[y, x]:
|
|
continue
|
|
visited[y, x] = True
|
|
if y > 0:
|
|
queue.append((y - 1, x))
|
|
if y + 1 < h:
|
|
queue.append((y + 1, x))
|
|
if x > 0:
|
|
queue.append((y, x - 1))
|
|
if x + 1 < w:
|
|
queue.append((y, x + 1))
|
|
|
|
return visited
|