forked from xiangwang25/whale-town-end-v2
Initial WhaleTown V2 backend
This commit is contained in:
293
scripts/skin_generation/tools/birefnet_cutout.py
Normal file
293
scripts/skin_generation/tools/birefnet_cutout.py
Normal file
@@ -0,0 +1,293 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Create clean transparent game-asset cutouts with BiRefNet.
|
||||
|
||||
This tool is intended for Novamailio-generated WhaleTown assets that come back
|
||||
on a plain matte background. It uses BiRefNet for the alpha mask, then fills
|
||||
transparent/semitransparent edge RGB from confident foreground pixels so Godot
|
||||
texture filtering cannot reveal the original matte color.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from collections import deque
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import Iterable, Tuple
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
DEFAULT_MODEL = "ZhengPeng7/BiRefNet"
|
||||
DEFAULT_HF_ENDPOINT = "https://hf-mirror.com"
|
||||
DEFAULT_SIZE = 1024
|
||||
|
||||
|
||||
def _die(message: str, code: int = 1) -> None:
|
||||
print(f"Error: {message}", file=sys.stderr)
|
||||
raise SystemExit(code)
|
||||
|
||||
|
||||
def _warn(message: str) -> None:
|
||||
print(f"Warning: {message}", file=sys.stderr)
|
||||
|
||||
|
||||
def _import_ml_deps() -> Tuple[object, object, object]:
|
||||
try:
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
from transformers import AutoModelForImageSegmentation
|
||||
except ImportError as exc:
|
||||
_die(
|
||||
"Missing BiRefNet dependencies. Install them with:\n"
|
||||
" python3 -m pip install torch torchvision transformers timm einops kornia scipy\n"
|
||||
f"Original import error: {exc}"
|
||||
)
|
||||
return torch, transforms, AutoModelForImageSegmentation
|
||||
|
||||
|
||||
def _select_device(torch: object, requested: str) -> str:
|
||||
if requested != "auto":
|
||||
return requested
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
return "cpu"
|
||||
|
||||
|
||||
def _load_birefnet(model_name: str, device: str, torch: object, auto_model: object) -> object:
|
||||
model = auto_model.from_pretrained(model_name, trust_remote_code=True)
|
||||
model.to(device)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def _run_birefnet(
|
||||
image: Image.Image,
|
||||
*,
|
||||
model_name: str,
|
||||
device: str,
|
||||
input_size: int,
|
||||
) -> Image.Image:
|
||||
torch, transforms, auto_model = _import_ml_deps()
|
||||
selected_device = _select_device(torch, device)
|
||||
print(f"BiRefNet device: {selected_device}", file=sys.stderr)
|
||||
print(f"BiRefNet model: {model_name}", file=sys.stderr)
|
||||
model = _load_birefnet(model_name, selected_device, torch, auto_model)
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(input_size, input_size),
|
||||
interpolation=transforms.InterpolationMode.BILINEAR,
|
||||
),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
]
|
||||
)
|
||||
input_tensor = transform(image).unsqueeze(0).to(selected_device)
|
||||
with torch.no_grad():
|
||||
output = model(input_tensor)
|
||||
if isinstance(output, (list, tuple)):
|
||||
prediction = output[-1]
|
||||
elif hasattr(output, "logits"):
|
||||
prediction = output.logits
|
||||
else:
|
||||
prediction = output
|
||||
prediction = prediction.sigmoid().detach().float().cpu()[0]
|
||||
if prediction.ndim == 3:
|
||||
prediction = prediction.squeeze(0)
|
||||
mask = transforms.ToPILImage()(prediction)
|
||||
return mask.resize(image.size, Image.Resampling.LANCZOS).convert("L")
|
||||
|
||||
|
||||
def _stabilize_mask(
|
||||
mask: Image.Image,
|
||||
*,
|
||||
low_cut: float,
|
||||
high_span: float,
|
||||
hard_low: float,
|
||||
hard_high: float,
|
||||
) -> Image.Image:
|
||||
mask_arr = np.asarray(mask).astype(np.float32) / 255.0
|
||||
mask_arr = np.clip((mask_arr - low_cut) / high_span, 0.0, 1.0)
|
||||
mask_arr = np.where(mask_arr > hard_high, 1.0, mask_arr)
|
||||
mask_arr = np.where(mask_arr < hard_low, 0.0, mask_arr)
|
||||
return Image.fromarray((mask_arr * 255).astype(np.uint8))
|
||||
|
||||
|
||||
def _connected_matte_mask(image: Image.Image, threshold: float) -> np.ndarray:
|
||||
"""Fallback mask for comparing or rescuing BiRefNet failures."""
|
||||
rgb = np.asarray(image.convert("RGB"))
|
||||
arr = rgb.astype(np.int32)
|
||||
h, w = arr.shape[:2]
|
||||
strips = np.concatenate(
|
||||
[
|
||||
arr[:80, :, :].reshape(-1, 3),
|
||||
arr[max(0, h - 80) : h, :, :].reshape(-1, 3),
|
||||
arr[:, :80, :].reshape(-1, 3),
|
||||
arr[:, max(0, w - 80) : w, :].reshape(-1, 3),
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
background = np.median(strips, axis=0).astype(np.int32)
|
||||
color_dist = np.sqrt(((arr - background) ** 2).sum(axis=2))
|
||||
candidate = color_dist <= threshold
|
||||
visited = np.zeros((h, w), dtype=bool)
|
||||
queue: deque[Tuple[int, int]] = deque()
|
||||
for x in range(w):
|
||||
for y in (0, h - 1):
|
||||
if candidate[y, x] and not visited[y, x]:
|
||||
visited[y, x] = True
|
||||
queue.append((y, x))
|
||||
for y in range(h):
|
||||
for x in (0, w - 1):
|
||||
if candidate[y, x] and not visited[y, x]:
|
||||
visited[y, x] = True
|
||||
queue.append((y, x))
|
||||
while queue:
|
||||
y, x = queue.popleft()
|
||||
for ny, nx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
|
||||
if 0 <= ny < h and 0 <= nx < w and (not visited[ny, nx]) and candidate[ny, nx]:
|
||||
visited[ny, nx] = True
|
||||
queue.append((ny, nx))
|
||||
return ~visited
|
||||
|
||||
|
||||
def _decontaminate_edge_rgb(rgb: np.ndarray, alpha: np.ndarray, confidence: int) -> np.ndarray:
|
||||
foreground = alpha > confidence
|
||||
if not foreground.any():
|
||||
_warn("Mask has no confident foreground pixels; edge decontamination skipped.")
|
||||
return rgb
|
||||
try:
|
||||
from scipy import ndimage
|
||||
except ImportError:
|
||||
_warn("scipy is missing; install scipy for edge RGB decontamination.")
|
||||
return rgb
|
||||
|
||||
_, indices = ndimage.distance_transform_edt(~foreground, return_indices=True)
|
||||
nearest_rgb = rgb[indices[0], indices[1]]
|
||||
alpha_f = alpha.astype(np.float32) / 255.0
|
||||
edge_mix = np.clip((0.98 - alpha_f) / 0.98, 0.0, 1.0)[..., None]
|
||||
replace_strength = np.where(alpha_f[..., None] < 0.98, edge_mix, 0.0)
|
||||
cleaned = rgb * (1.0 - replace_strength) + nearest_rgb * replace_strength
|
||||
return np.clip(cleaned, 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def _compose_cutout(image: Image.Image, mask: Image.Image, confidence: int) -> Image.Image:
|
||||
rgb = np.asarray(image.convert("RGB")).astype(np.float32)
|
||||
alpha = np.asarray(mask).astype(np.uint8)
|
||||
cleaned_rgb = _decontaminate_edge_rgb(rgb, alpha, confidence)
|
||||
return Image.fromarray(np.dstack([cleaned_rgb, alpha]))
|
||||
|
||||
|
||||
def _make_preview(cutout: Image.Image, output: Path, scale: float) -> None:
|
||||
w, h = cutout.size
|
||||
small = cutout.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
|
||||
canvas = Image.new("RGB", (small.width * 2 + 48, small.height + 48), (238, 238, 238))
|
||||
|
||||
magenta_plate = Image.new("RGB", small.size, (198, 76, 190))
|
||||
magenta_plate.paste(small, (0, 0), small)
|
||||
|
||||
checker = Image.new("RGB", small.size, (230, 230, 230))
|
||||
draw = ImageDraw.Draw(checker)
|
||||
step = 32
|
||||
for y in range(0, small.height, step):
|
||||
for x in range(0, small.width, step):
|
||||
if ((x // step) + (y // step)) % 2 == 0:
|
||||
draw.rectangle([x, y, x + step - 1, y + step - 1], fill=(190, 206, 224))
|
||||
checker.paste(small, (0, 0), small)
|
||||
|
||||
canvas.paste(magenta_plate, (16, 32))
|
||||
canvas.paste(checker, (small.width + 32, 32))
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
canvas.save(output)
|
||||
|
||||
|
||||
def _magenta_stats(image: Image.Image) -> Tuple[int, int, int]:
|
||||
arr = np.asarray(image.convert("RGBA")).astype(np.int16)
|
||||
rgb = arr[:, :, :3]
|
||||
alpha = arr[:, :, 3]
|
||||
semi = (alpha > 0) & (alpha < 255)
|
||||
magenta = (alpha > 0) & (rgb[:, :, 0] > 120) & (rgb[:, :, 1] < 60) & (rgb[:, :, 2] > 90)
|
||||
semi_magenta = semi & magenta
|
||||
return int(semi.sum()), int(magenta.sum()), int(semi_magenta.sum())
|
||||
|
||||
|
||||
def parse_args(argv: Iterable[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--input", required=True, help="Source RGB/RGBA image with matte background.")
|
||||
parser.add_argument("--output", required=True, help="Transparent PNG output path.")
|
||||
parser.add_argument("--mask-out", help="Optional alpha mask output path.")
|
||||
parser.add_argument("--preview-out", help="Optional preview sheet output path.")
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help=f"HF model id. Default: {DEFAULT_MODEL}")
|
||||
parser.add_argument("--hf-endpoint", default=DEFAULT_HF_ENDPOINT, help="Hugging Face endpoint/mirror.")
|
||||
parser.add_argument("--device", choices=["auto", "cpu", "mps", "cuda"], default="auto")
|
||||
parser.add_argument("--input-size", type=int, default=DEFAULT_SIZE, help="Square BiRefNet input size.")
|
||||
parser.add_argument("--low-cut", type=float, default=0.025, help="Low alpha normalization cut.")
|
||||
parser.add_argument("--high-span", type=float, default=0.94, help="Alpha normalization span.")
|
||||
parser.add_argument("--hard-low", type=float, default=0.015, help="Values below this become transparent.")
|
||||
parser.add_argument("--hard-high", type=float, default=0.985, help="Values above this become opaque.")
|
||||
parser.add_argument("--edge-confidence", type=int, default=245, help="Confident foreground alpha for RGB fill.")
|
||||
parser.add_argument("--preview-scale", type=float, default=0.52)
|
||||
parser.add_argument("--fallback-connected-matte", action="store_true", help="Use simple connected matte mask instead of BiRefNet.")
|
||||
parser.add_argument("--fallback-threshold", type=float, default=28.0)
|
||||
return parser.parse_args(list(argv))
|
||||
|
||||
|
||||
def main(argv: Iterable[str]) -> int:
|
||||
args = parse_args(argv)
|
||||
input_path = Path(args.input)
|
||||
output_path = Path(args.output)
|
||||
if not input_path.exists():
|
||||
_die(f"Input image not found: {input_path}")
|
||||
|
||||
os.environ.setdefault("HF_ENDPOINT", args.hf_endpoint)
|
||||
image = Image.open(input_path).convert("RGB")
|
||||
if args.fallback_connected_matte:
|
||||
foreground = _connected_matte_mask(image, args.fallback_threshold)
|
||||
mask = Image.fromarray(foreground.astype(np.uint8) * 255)
|
||||
else:
|
||||
mask = _run_birefnet(
|
||||
image,
|
||||
model_name=args.model,
|
||||
device=args.device,
|
||||
input_size=args.input_size,
|
||||
)
|
||||
mask = _stabilize_mask(
|
||||
mask,
|
||||
low_cut=args.low_cut,
|
||||
high_span=args.high_span,
|
||||
hard_low=args.hard_low,
|
||||
hard_high=args.hard_high,
|
||||
)
|
||||
|
||||
cutout = _compose_cutout(image, mask, args.edge_confidence)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cutout.save(output_path)
|
||||
|
||||
if args.mask_out:
|
||||
mask_path = Path(args.mask_out)
|
||||
mask_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
mask.save(mask_path)
|
||||
if args.preview_out:
|
||||
_make_preview(cutout, Path(args.preview_out), args.preview_scale)
|
||||
|
||||
semi, magenta, semi_magenta = _magenta_stats(cutout)
|
||||
print(f"wrote {output_path}")
|
||||
if args.mask_out:
|
||||
print(f"wrote {args.mask_out}")
|
||||
if args.preview_out:
|
||||
print(f"wrote {args.preview_out}")
|
||||
print(f"alpha_bbox={cutout.getchannel('A').getbbox()}")
|
||||
print(f"semi_transparent_pixels={semi}")
|
||||
print(f"visible_magenta_like_pixels={magenta}")
|
||||
print(f"semi_magenta_like_pixels={semi_magenta}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(sys.argv[1:]))
|
||||
Reference in New Issue
Block a user