Files
AI_spine/normalize_reference.py
T
2026-09-17 18:50:43 +08:00

60 lines
2.4 KiB
Python

#!/usr/bin/env python3
"""Prepare a reference for feature matching without changing its artwork."""
import argparse
from pathlib import Path
from PIL import Image
import numpy as np
import cv2
def prepare_rgba(image, white_threshold=240):
"""Keep existing alpha; remove only border-connected near-white background.
No crop or rescale here: semantic coordinates must retain the full canvas.
White regions enclosed by character outlines are deliberately preserved.
"""
rgba = np.array(image.convert('RGBA'))
alpha = rgba[:, :, 3]
method = 'existing_alpha'
removed = 0
if not np.any(alpha == 0):
white = np.all(rgba[:, :, :3] >= white_threshold, axis=2).astype(np.uint8)
_, labels = cv2.connectedComponents(white, connectivity=4)
border = np.unique(np.concatenate((labels[0], labels[-1], labels[:, 0], labels[:, -1])))
background = np.isin(labels, border[border != 0])
removed = int(background.sum())
rgba[background, 3] = 0
method = 'border_connected_white'
foreground = rgba[:, :, 3] > 0
if not foreground.any():
raise ValueError('Reference has no foreground after background removal')
if foreground.all():
raise ValueError('No transparent background or border-connected white background; supply a cutout RGBA PNG')
return Image.fromarray(rgba), {
'method': method, 'white_threshold': white_threshold,
'removed_pixels': removed, 'foreground_pixels': int(foreground.sum()),
'width': image.width, 'height': image.height,
'needs_review': method != 'existing_alpha',
}
def main():
p = argparse.ArgumentParser()
p.add_argument('--input', required=True)
p.add_argument('--output', required=True)
p.add_argument('--target', type=int, default=1024)
p.add_argument('--remove-background', action='store_true')
a = p.parse_args()
im = Image.open(a.input).convert('RGBA')
if a.remove_background:
im, _ = prepare_rgba(im)
longest = max(im.width, im.height)
if longest < a.target:
scale = a.target / longest
size = (max(1, round(im.width * scale)), max(1, round(im.height * scale)))
# Nearest-neighbour preserves pixel-art edges and transparency.
im = im.resize(size, Image.Resampling.NEAREST)
im.save(a.output)
print(f'Reference prepared: {im.width}x{im.height}')
if __name__ == '__main__': main()