87 lines
3.7 KiB
Python
87 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Deterministically split a GPT generated white-background parts sheet.
|
|
|
|
The model call is intentionally kept in Genarrative's Node runner. This file
|
|
only performs the open-source, deterministic connected-components step used by
|
|
the Genielabs Spine Animation AI workflow.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
|
|
try:
|
|
import cv2
|
|
except ImportError as error:
|
|
raise SystemExit("缺少 OpenCV，请先执行：python3 -m pip install opencv-python") from error
|
|
|
|
|
|
def segment(atlas_path: str, output_dir: str, min_area: int, padding: int, bg_threshold: int):
|
|
image = cv2.imread(atlas_path, cv2.IMREAD_UNCHANGED)
|
|
if image is None:
|
|
raise ValueError(f"无法读取拆件图集：{atlas_path}")
|
|
if image.ndim == 2:
|
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGRA)
|
|
elif image.shape[2] == 3:
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
|
|
|
|
# Removing every bright pixel also erased eye whites, teeth and pale
|
|
# clothing. Treat white connected to the canvas boundary as background;
|
|
# enclosed highlights remain part of their surrounding cutout.
|
|
white = np.all(image[:, :, :3] >= bg_threshold, axis=2)
|
|
transparent = image[:, :, 3] == 0
|
|
bg_candidates = (white | transparent).astype(np.uint8)
|
|
_, bg_labels, _, _ = cv2.connectedComponentsWithStats(bg_candidates, 8)
|
|
border_labels = np.unique(np.concatenate((bg_labels[0], bg_labels[-1], bg_labels[:, 0], bg_labels[:, -1])))
|
|
border_labels = border_labels[border_labels != 0]
|
|
background = np.isin(bg_labels, border_labels) | transparent
|
|
mask = (~background).astype(np.uint8) * 255
|
|
count, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
height, width = image.shape[:2]
|
|
saved = []
|
|
components = []
|
|
for label in range(1, count):
|
|
area = int(stats[label, cv2.CC_STAT_AREA])
|
|
if area < min_area:
|
|
continue
|
|
x, y = int(stats[label, cv2.CC_STAT_LEFT]), int(stats[label, cv2.CC_STAT_TOP])
|
|
w, h = int(stats[label, cv2.CC_STAT_WIDTH]), int(stats[label, cv2.CC_STAT_HEIGHT])
|
|
x1, y1 = max(0, x - padding), max(0, y - padding)
|
|
x2, y2 = min(width, x + w + padding), min(height, y + h + padding)
|
|
crop = image[y1:y2, x1:x2].copy()
|
|
crop[labels[y1:y2, x1:x2] != label] = (0, 0, 0, 0)
|
|
name = f"part_{len(saved):02d}.png"
|
|
path = os.path.join(output_dir, name)
|
|
if not cv2.imwrite(path, crop):
|
|
raise ValueError(f"无法写入拆件：{path}")
|
|
saved.append(path)
|
|
components.append({"name": name[:-4], "x": x, "y": y, "width": w, "height": h, "area": area})
|
|
if not saved:
|
|
raise ValueError("拆件图集没有检测到前景，请检查白底和 bg-threshold")
|
|
with open(os.path.join(output_dir, "parts-manifest.json"), "w", encoding="utf-8") as handle:
|
|
json.dump({"atlas_width": width, "atlas_height": height, "parts": components}, handle, ensure_ascii=False, indent=2)
|
|
return saved
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--atlas", required=True)
|
|
parser.add_argument("--output-dir", required=True)
|
|
parser.add_argument("--min-area", type=int, default=500)
|
|
parser.add_argument("--padding", type=int, default=12)
|
|
parser.add_argument("--bg-threshold", type=int, default=240)
|
|
args = parser.parse_args()
|
|
try:
|
|
paths = segment(args.atlas, args.output_dir, args.min_area, args.padding, args.bg_threshold)
|
|
except Exception as error:
|
|
print(f"ERROR: {error}", file=sys.stderr)
|
|
return 1
|
|
print(f"拆出 {len(paths)} 个部件")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|