306 lines
11 KiB
Python
306 lines
11 KiB
Python
import os
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageEnhance
|
|
|
|
from .palettes import PALETTE_NAMES, get_palette_rgb
|
|
|
|
# The compiled engine lives next to this file, inside the custom node folder, so it
|
|
# works standalone on any Windows machine. To use your own build, just replace this
|
|
# .exe in the node folder.
|
|
_HERE = os.path.dirname(os.path.abspath(__file__))
|
|
DEFAULT_EXE = os.path.join(_HERE, "spritefusion-pixel-snapper.exe")
|
|
|
|
|
|
def tensor_to_pil(img_tensor):
|
|
arr = (img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
return Image.fromarray(arr, mode="RGB")
|
|
|
|
|
|
def pil_to_tensor(img):
|
|
arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr)
|
|
|
|
|
|
def remap_to_palette(img, rgb_palette, dither):
|
|
"""Reduce the image to the exact colors of a fixed palette."""
|
|
pal_img = Image.new("P", (1, 1))
|
|
flat = []
|
|
for r, g, b in rgb_palette:
|
|
flat.extend([r, g, b])
|
|
flat.extend([0, 0, 0] * (256 - len(rgb_palette)))
|
|
pal_img.putpalette(flat)
|
|
dither_mode = Image.FLOYDSTEINBERG if dither else Image.NONE
|
|
return img.convert("RGB").quantize(
|
|
palette=pal_img, colors=len(rgb_palette), dither=dither_mode
|
|
).convert("RGB")
|
|
|
|
|
|
def correct_aspect(img, src_size):
|
|
"""Fix the 'flattening' from the engine: the edge-walker can emit a slightly
|
|
different number of cells on X vs Y, distorting the aspect ratio. We trim the
|
|
over-represented axis (nearest) so the output aspect matches the source."""
|
|
ow, oh = img.size
|
|
sw, sh = src_size
|
|
if min(ow, oh, sw, sh) <= 0:
|
|
return img
|
|
src_aspect = sw / sh
|
|
out_aspect = ow / oh
|
|
if abs(out_aspect - src_aspect) < 1e-3:
|
|
return img
|
|
if out_aspect > src_aspect:
|
|
new_w, new_h = max(1, round(oh * src_aspect)), oh # too wide -> trim width
|
|
else:
|
|
new_w, new_h = ow, max(1, round(ow / src_aspect)) # too tall -> trim height
|
|
if (new_w, new_h) == (ow, oh):
|
|
return img
|
|
return img.resize((new_w, new_h), Image.NEAREST)
|
|
|
|
|
|
def run_snapper(exe_path, pil_in, k):
|
|
"""Save the image, run the spritefusion-pixel-snapper binary and read the result."""
|
|
if not (isinstance(exe_path, str) and os.path.isfile(exe_path)):
|
|
raise FileNotFoundError(
|
|
f"Pixel Snapper engine not found. Place 'spritefusion-pixel-snapper.exe' "
|
|
f"in the node folder: {_HERE}"
|
|
)
|
|
tmpdir = tempfile.mkdtemp(prefix="pixelsnap_")
|
|
in_path = os.path.join(tmpdir, "in.png")
|
|
out_path = os.path.join(tmpdir, "out.png")
|
|
pil_in.save(in_path)
|
|
result = subprocess.run(
|
|
[exe_path, in_path, out_path, str(int(k))],
|
|
capture_output=True, text=True,
|
|
)
|
|
if result.returncode != 0 or not os.path.isfile(out_path):
|
|
raise RuntimeError(
|
|
f"Pixel Snapper failed (code {result.returncode}).\n"
|
|
f"stdout: {result.stdout}\nstderr: {result.stderr}"
|
|
)
|
|
out_img = Image.open(out_path).convert("RGB").copy()
|
|
return out_img
|
|
|
|
|
|
class PixelSnapperNode:
|
|
"""Calls the compiled engine 'spritefusion-pixel-snapper.exe' which detects the
|
|
native pixel art grid and reduces the image to k colors. Supports color
|
|
adjustments (pre) and remapping to a fixed video game palette (post)."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"colors_k": ("INT", {"default": 16, "min": 2, "max": 256, "step": 1}),
|
|
"saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
|
|
"contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
|
|
"brightness": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
|
|
"fix_aspect": ("BOOLEAN", {"default": True}),
|
|
"keep_output_size": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"palette": ("GAME_PALETTE",),
|
|
"palette_dither": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "image/pixel_art"
|
|
|
|
def run(self, image, colors_k, saturation, contrast, brightness,
|
|
fix_aspect, keep_output_size, palette=None, palette_dither=False):
|
|
# If a palette is connected, the k passed to the engine is the palette size.
|
|
k = len(palette) if palette is not None else colors_k
|
|
|
|
out_batch = []
|
|
target_size = None
|
|
for i in range(image.shape[0]):
|
|
img = tensor_to_pil(image[i])
|
|
orig_size = img.size
|
|
|
|
# Color adjustments BEFORE the snap (they affect the engine's k-means).
|
|
if brightness != 1.0:
|
|
img = ImageEnhance.Brightness(img).enhance(brightness)
|
|
if contrast != 1.0:
|
|
img = ImageEnhance.Contrast(img).enhance(contrast)
|
|
if saturation != 1.0:
|
|
img = ImageEnhance.Color(img).enhance(saturation)
|
|
|
|
result = run_snapper(DEFAULT_EXE, img, k)
|
|
|
|
# Fix the engine's aspect-ratio flattening (edge-walker cell mismatch).
|
|
if fix_aspect:
|
|
result = correct_aspect(result, orig_size)
|
|
|
|
# Remap to a fixed palette AFTER the snap (exact video game colors).
|
|
if palette is not None:
|
|
result = remap_to_palette(result, palette, palette_dither)
|
|
|
|
if keep_output_size:
|
|
result = result.resize(orig_size, Image.NEAREST)
|
|
|
|
# Ensure consistent size across the batch (torch.stack requires it).
|
|
if target_size is None:
|
|
target_size = result.size
|
|
elif result.size != target_size:
|
|
result = result.resize(target_size, Image.NEAREST)
|
|
|
|
out_batch.append(pil_to_tensor(result))
|
|
|
|
return (torch.stack(out_batch),)
|
|
|
|
|
|
class VideoGamePaletteNode:
|
|
"""Provides a classic video game color palette for the Pixel Snapper."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"palette_name": (PALETTE_NAMES, {"default": PALETTE_NAMES[0]}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("GAME_PALETTE",)
|
|
RETURN_NAMES = ("palette",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "image/pixel_art"
|
|
|
|
def run(self, palette_name):
|
|
return (get_palette_rgb(palette_name),)
|
|
|
|
|
|
RESAMPLE = {
|
|
"nearest": Image.NEAREST,
|
|
"box": Image.BOX,
|
|
"bilinear": Image.BILINEAR,
|
|
"hamming": Image.HAMMING,
|
|
"bicubic": Image.BICUBIC,
|
|
"lanczos": Image.LANCZOS,
|
|
}
|
|
|
|
|
|
class PixelArtResizerNode:
|
|
"""Resizes pixel art. To scale up without blurring use 'nearest'; to shrink a
|
|
large image down to pixel art resolution use 'box' or 'nearest'."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mode": (["scale_factor", "target_size"], {"default": "scale_factor"}),
|
|
"scale_factor": ("FLOAT", {"default": 2.0, "min": 0.05, "max": 64.0, "step": 0.05}),
|
|
"width": ("INT", {"default": 128, "min": 1, "max": 8192, "step": 1}),
|
|
"height": ("INT", {"default": 128, "min": 1, "max": 8192, "step": 1}),
|
|
"interpolation": (list(RESAMPLE.keys()), {"default": "nearest"}),
|
|
"keep_aspect": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "image/pixel_art"
|
|
|
|
def run(self, image, mode, scale_factor, width, height, interpolation, keep_aspect):
|
|
filt = RESAMPLE[interpolation]
|
|
out_batch = []
|
|
target = None
|
|
for i in range(image.shape[0]):
|
|
img = tensor_to_pil(image[i])
|
|
w, h = img.size
|
|
|
|
if mode == "scale_factor":
|
|
nw = max(1, round(w * scale_factor))
|
|
nh = max(1, round(h * scale_factor))
|
|
else:
|
|
if keep_aspect:
|
|
ratio = min(width / w, height / h)
|
|
nw = max(1, round(w * ratio))
|
|
nh = max(1, round(h * ratio))
|
|
else:
|
|
nw, nh = int(width), int(height)
|
|
|
|
result = img.resize((nw, nh), filt)
|
|
|
|
# Batch needs a consistent size for torch.stack.
|
|
if target is None:
|
|
target = result.size
|
|
elif result.size != target:
|
|
result = result.resize(target, filt)
|
|
|
|
out_batch.append(pil_to_tensor(result))
|
|
|
|
return (torch.stack(out_batch),)
|
|
|
|
|
|
class PixelArtPreviewNode:
|
|
"""Pixel art preview with crisp rendering (nearest-neighbor) and an interactive
|
|
magnifier. Hover over the image to zoom in. Output node."""
|
|
|
|
def __init__(self):
|
|
self._prefix = "pixelart_preview"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"magnifier": ("BOOLEAN", {"default": True}),
|
|
"zoom": ("INT", {"default": 6, "min": 2, "max": 32, "step": 1}),
|
|
"lens_size": ("INT", {"default": 140, "min": 40, "max": 400, "step": 10}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "preview"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image/pixel_art"
|
|
|
|
def preview(self, images, magnifier, zoom, lens_size, prompt=None, extra_pnginfo=None):
|
|
import folder_paths
|
|
|
|
output_dir = folder_paths.get_temp_directory()
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
|
self._prefix, output_dir, images[0].shape[1], images[0].shape[0]
|
|
)
|
|
results = []
|
|
for image in images:
|
|
arr = (255.0 * image.cpu().numpy()).clip(0, 255).astype(np.uint8)
|
|
img = Image.fromarray(arr)
|
|
file = f"{filename}_{counter:05}_.png"
|
|
img.save(os.path.join(full_output_folder, file), compress_level=1)
|
|
results.append({"filename": file, "subfolder": subfolder, "type": "temp"})
|
|
counter += 1
|
|
|
|
return {
|
|
"ui": {
|
|
"images": results,
|
|
"magnifier": [magnifier],
|
|
"zoom": [zoom],
|
|
"lens_size": [lens_size],
|
|
}
|
|
}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PixelSnapper": PixelSnapperNode,
|
|
"VideoGamePalette": VideoGamePaletteNode,
|
|
"PixelArtResizer": PixelArtResizerNode,
|
|
"PixelArtPreview": PixelArtPreviewNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PixelSnapper": "Pixel Snapper (Engine)",
|
|
"VideoGamePalette": "Video Game Palette",
|
|
"PixelArtResizer": "Pixel Art Resizer",
|
|
"PixelArtPreview": "Pixel Art Preview",
|
|
}
|