Improve median fixer algo
This commit is contained in:
+4
-1
@@ -4,7 +4,8 @@ from .nodes import (
|
||||
PixelGrid_MergeSimilar,
|
||||
PixelGrid_Analyze,
|
||||
PixelGrid_PaletteToImage,
|
||||
GridMedianFixer
|
||||
GridMedianFixer,
|
||||
PixelGrid_CreateSolidColorImage
|
||||
)
|
||||
from .autofixer import PixelGrid_Autofixer
|
||||
from .color_names import PixelGrid_HexToNames
|
||||
@@ -16,6 +17,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"PixelGrid_Analyze": PixelGrid_Analyze,
|
||||
"PixelGrid_PaletteToImage": PixelGrid_PaletteToImage,
|
||||
"GridMedianFixer": GridMedianFixer,
|
||||
"PixelGrid_CreateSolidColorImage": PixelGrid_CreateSolidColorImage,
|
||||
"PixelGrid_Autofixer": PixelGrid_Autofixer,
|
||||
"PixelGrid_HexToNames": PixelGrid_HexToNames
|
||||
}
|
||||
@@ -27,6 +29,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PixelGrid_Analyze": "Pixel Grid: Analyze Palette",
|
||||
"PixelGrid_PaletteToImage": "Pixel Grid: Palette to Image",
|
||||
"GridMedianFixer": "Pixel Grid: Median Fixer",
|
||||
"PixelGrid_CreateSolidColorImage": "Pixel Grid: Create Solid Color Image",
|
||||
"PixelGrid_Autofixer": "Pixel Grid: Autofixer",
|
||||
"PixelGrid_HexToNames": "Pixel Grid: HEX to Color Names"
|
||||
}
|
||||
|
||||
@@ -8,6 +8,37 @@ def hex_to_rgb(hex_str):
|
||||
def rgb_to_hex(rgb):
|
||||
return '#{:02x}{:02x}{:02x}'.format(int(rgb[0]), int(rgb[1]), int(rgb[2]))
|
||||
|
||||
class PixelGrid_CreateSolidColorImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", {"default": 512, "min": 1, "max": 8192}),
|
||||
"height": ("INT", {"default": 512, "min": 1, "max": 8192}),
|
||||
"hex_code": ("STRING", {"default": "#FF0000"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "create_image"
|
||||
CATEGORY = "Pixel Grid Helpers"
|
||||
|
||||
def create_image(self, width, height, hex_code):
|
||||
match = re.fullmatch(r"#?([0-9a-fA-F]{3}|[0-9a-fA-F]{6})", hex_code.strip())
|
||||
if match is None:
|
||||
raise ValueError(
|
||||
f"Invalid HEX color {hex_code!r}; use 3 or 6 hexadecimal digits."
|
||||
)
|
||||
|
||||
digits = match.group(1)
|
||||
if len(digits) == 3:
|
||||
digits = "".join(digit * 2 for digit in digits)
|
||||
rgb = [int(digits[index:index + 2], 16) for index in (0, 2, 4)]
|
||||
color = torch.tensor(rgb, dtype=torch.float32).div_(255).view(1, 1, 1, 3)
|
||||
image = color.expand(1, height, width, 3).clone()
|
||||
return (image,)
|
||||
|
||||
class PixelGrid_KMeans:
|
||||
"""Node 1: Enforce Max N Colors (K-Means Quantization)"""
|
||||
@classmethod
|
||||
@@ -229,7 +260,8 @@ class PixelGrid_PaletteToImage:
|
||||
|
||||
# Re-including the previous node requested
|
||||
class GridMedianFixer:
|
||||
"""Previous Node: Median Filter for Logic Grids"""
|
||||
"""Median-reduce an image using the most coherent pixel-grid phase."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
@@ -238,21 +270,134 @@ class GridMedianFixer:
|
||||
"grid_size": ("INT", {"default": 6, "min": 1, "max": 128}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("downscaled_image", "original_size_image")
|
||||
FUNCTION = "process_grid"
|
||||
CATEGORY = "Pixel Grid Helpers"
|
||||
|
||||
@staticmethod
|
||||
def _phase_downscaled(image, grid_size, y_offset, x_offset):
|
||||
batch, height, width, channels = image.shape
|
||||
rows = max((height - y_offset) // grid_size, 0)
|
||||
columns = max((width - x_offset) // grid_size, 0)
|
||||
end_y = y_offset + rows * grid_size
|
||||
end_x = x_offset + columns * grid_size
|
||||
cropped = image[:, y_offset:end_y, x_offset:end_x, :]
|
||||
|
||||
blocks = cropped.reshape(
|
||||
batch, rows, grid_size, columns, grid_size, channels
|
||||
).permute(0, 1, 3, 2, 4, 5)
|
||||
blocks = blocks.reshape(
|
||||
batch, rows, columns, grid_size * grid_size, channels
|
||||
)
|
||||
return torch.median(blocks, dim=3).values
|
||||
|
||||
@classmethod
|
||||
def _phase_outputs(cls, image, grid_size, y_offset, x_offset):
|
||||
downscaled = cls._phase_downscaled(
|
||||
image, grid_size, y_offset, x_offset
|
||||
)
|
||||
upscaled = downscaled.repeat_interleave(grid_size, dim=1).repeat_interleave(
|
||||
grid_size, dim=2
|
||||
)
|
||||
return downscaled, upscaled
|
||||
|
||||
@staticmethod
|
||||
def _common_phase_bounds(length, grid_size):
|
||||
valid_offsets = [
|
||||
offset
|
||||
for offset in range(grid_size)
|
||||
if length - offset >= grid_size
|
||||
]
|
||||
if not valid_offsets:
|
||||
return 0, 0
|
||||
|
||||
start = valid_offsets[-1]
|
||||
end = min(
|
||||
offset + ((length - offset) // grid_size) * grid_size
|
||||
for offset in valid_offsets
|
||||
)
|
||||
return start, end
|
||||
|
||||
@staticmethod
|
||||
def _phase_error(
|
||||
image,
|
||||
grid_size,
|
||||
y_offset,
|
||||
x_offset,
|
||||
score_top,
|
||||
score_bottom,
|
||||
score_left,
|
||||
score_right,
|
||||
):
|
||||
batch, _, _, channels = image.shape
|
||||
first_row = (score_top - y_offset) // grid_size
|
||||
last_row = (score_bottom - 1 - y_offset) // grid_size + 1
|
||||
first_column = (score_left - x_offset) // grid_size
|
||||
last_column = (score_right - 1 - x_offset) // grid_size + 1
|
||||
|
||||
block_top = y_offset + first_row * grid_size
|
||||
block_bottom = y_offset + last_row * grid_size
|
||||
block_left = x_offset + first_column * grid_size
|
||||
block_right = x_offset + last_column * grid_size
|
||||
rows = last_row - first_row
|
||||
columns = last_column - first_column
|
||||
score_pixels = image[:, block_top:block_bottom, block_left:block_right, :]
|
||||
blocks = score_pixels.reshape(
|
||||
batch, rows, grid_size, columns, grid_size, channels
|
||||
).permute(0, 1, 3, 2, 4, 5)
|
||||
blocks = blocks.reshape(
|
||||
batch, rows, columns, grid_size * grid_size, channels
|
||||
)
|
||||
score_medians = torch.median(blocks, dim=3).values
|
||||
aligned_region = score_medians.repeat_interleave(
|
||||
grid_size, dim=1
|
||||
).repeat_interleave(grid_size, dim=2)
|
||||
|
||||
local_top = score_top - block_top
|
||||
local_left = score_left - block_left
|
||||
region_height = score_bottom - score_top
|
||||
region_width = score_right - score_left
|
||||
aligned_region = aligned_region[
|
||||
:, local_top:local_top + region_height,
|
||||
local_left:local_left + region_width, :
|
||||
]
|
||||
source_region = image[:, score_top:score_bottom, score_left:score_right, :]
|
||||
error = source_region.to(torch.float32) - aligned_region.to(torch.float32)
|
||||
return error.abs().mean()
|
||||
|
||||
def process_grid(self, image, grid_size):
|
||||
B, H, W, C = image.shape
|
||||
h_mod, w_mod = H % grid_size, W % grid_size
|
||||
if h_mod > 0 or w_mod > 0:
|
||||
image = image[:, :H-h_mod, :W-w_mod, :]
|
||||
B, H, W, C = image.shape
|
||||
|
||||
reshaped = image.view(B, H // grid_size, grid_size, W // grid_size, grid_size, C)
|
||||
permuted = reshaped.permute(0, 1, 3, 2, 4, 5)
|
||||
flattened = permuted.reshape(B, H // grid_size, W // grid_size, grid_size * grid_size, C)
|
||||
downscaled, _ = torch.median(flattened, dim=3)
|
||||
upscaled = downscaled.repeat_interleave(grid_size, dim=1).repeat_interleave(grid_size, dim=2)
|
||||
return (downscaled, upscaled)
|
||||
if grid_size < 1:
|
||||
raise ValueError("grid_size must be at least 1")
|
||||
|
||||
_, height, width, _ = image.shape
|
||||
with torch.no_grad():
|
||||
if height < grid_size or width < grid_size:
|
||||
return self._phase_outputs(image, grid_size, 0, 0)
|
||||
|
||||
score_top, score_bottom = self._common_phase_bounds(height, grid_size)
|
||||
score_left, score_right = self._common_phase_bounds(width, grid_size)
|
||||
phases = []
|
||||
scores = []
|
||||
|
||||
y_offsets = range(min(grid_size, height - grid_size + 1))
|
||||
x_offsets = range(min(grid_size, width - grid_size + 1))
|
||||
for y_offset in y_offsets:
|
||||
for x_offset in x_offsets:
|
||||
score = self._phase_error(
|
||||
image,
|
||||
grid_size,
|
||||
y_offset,
|
||||
x_offset,
|
||||
score_top,
|
||||
score_bottom,
|
||||
score_left,
|
||||
score_right,
|
||||
)
|
||||
phases.append((y_offset, x_offset))
|
||||
scores.append(score)
|
||||
|
||||
best_index = torch.argmin(torch.stack(scores)).item()
|
||||
best_y, best_x = phases[best_index]
|
||||
return self._phase_outputs(image, grid_size, best_y, best_x)
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-pixelgrid-helpers"
|
||||
description = "A suite of nodes for pixel art: color quantization, palette enforcement, median filtering and grid logic."
|
||||
version = "1.1.0"
|
||||
version = "1.1.1"
|
||||
license = { text = "MIT" }
|
||||
dependencies = []
|
||||
|
||||
|
||||
@@ -136,9 +136,20 @@ to Preview Image or Save Image.
|
||||
Clean up enlarged pixel art with noisy or uneven pixel blocks. Use this when
|
||||
each logical pixel occupies a known square block, such as 6 by 6 image pixels.
|
||||
|
||||
Set `grid_size` to the block size. Use `downscaled_image` for native-resolution
|
||||
editing, or `original_size_image` to keep the enlarged appearance with solid blocks.
|
||||
Edges that do not fit a complete block are cropped.
|
||||
Set `grid_size` to the block size. The node tests each horizontal and vertical
|
||||
grid offset that leaves at least one complete block, then chooses the alignment
|
||||
with the lowest mean absolute RGB error from each pixel to its block median.
|
||||
This favors alignments where pixels within each logical block are most
|
||||
consistent; equal scores prefer the top-left alignment. Scoring uses the region
|
||||
covered by all candidate alignments, so incomplete edge fragments do not bias
|
||||
the comparison.
|
||||
|
||||
The selected leading offset and any incomplete trailing blocks are cropped.
|
||||
Use `downscaled_image` for native-resolution editing, or `original_size_image`
|
||||
to keep the selected grid enlarged as solid blocks. Output dimensions can differ
|
||||
from the input. The exhaustive phase search costs `grid_size` squared candidate
|
||||
alignments, so larger block sizes take longer; scoring runs on the input device,
|
||||
including the GPU when available.
|
||||
|
||||
| Input | Usage |
|
||||
| --- | --- |
|
||||
@@ -208,12 +219,28 @@ Different HEX colors can receive the same name; each name appears only once.
|
||||
Colors without an exact named match receive the closest available name.
|
||||
Whitespace and empty entries are ignored; malformed HEX colors raise an error.
|
||||
|
||||
### 9. Pixel Grid: Create Solid Color Image
|
||||
|
||||
Create an image filled with one HEX color. Set the output dimensions and enter
|
||||
a three- or six-digit HEX color, with or without the leading `#`.
|
||||
|
||||
| Input | Usage |
|
||||
| --- | --- |
|
||||
| `width` | Output image width in pixels. |
|
||||
| `height` | Output image height in pixels. |
|
||||
| `hex_code` | HEX color to fill the image with; default: `#FF0000`. |
|
||||
|
||||
| Output | Usage |
|
||||
| --- | --- |
|
||||
| `image` | The generated solid-color image. |
|
||||
|
||||
## Example Workflows
|
||||
|
||||
- **Clean generated pixel art:** Image -> Autofixer -> Preview Image or Save Image.
|
||||
- **Reuse a reference palette:** Reference image -> Analyze Palette -> Enforce Palette's `hex_palette`; connect your artwork to Enforce Palette's `image`.
|
||||
- **Inspect and describe colors:** Connect any `hex_palette` output to Palette to Image for swatches and HEX to Color Names for readable names.
|
||||
- **Recover a known pixel grid:** Image -> Median Fixer -> use `downscaled_image` for editing or `original_size_image` for an enlarged result.
|
||||
- **Create a flat background:** Create Solid Color Image -> Preview Image or combine it with another image.
|
||||
|
||||
## Attribution
|
||||
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
import importlib.util
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
MODULE_PATH = Path(__file__).parents[1] / "nodes.py"
|
||||
SPEC = importlib.util.spec_from_file_location("pixelgrid_nodes", MODULE_PATH)
|
||||
NODES = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(NODES)
|
||||
|
||||
|
||||
class CreateSolidColorImageTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = NODES.PixelGrid_CreateSolidColorImage()
|
||||
|
||||
def test_node_contract(self):
|
||||
self.assertEqual(self.node.RETURN_TYPES, ("IMAGE",))
|
||||
self.assertEqual(self.node.RETURN_NAMES, ("image",))
|
||||
self.assertEqual(self.node.CATEGORY, "Pixel Grid Helpers")
|
||||
self.assertEqual(self.node.FUNCTION, "create_image")
|
||||
self.assertEqual(
|
||||
set(self.node.INPUT_TYPES()["required"]),
|
||||
{"width", "height", "hex_code"},
|
||||
)
|
||||
|
||||
def test_creates_rgb_image_with_requested_dimensions_and_color(self):
|
||||
for hex_code, expected in (
|
||||
("#123456", [0x12 / 255, 0x34 / 255, 0x56 / 255]),
|
||||
(" f0A ", [1.0, 0.0, 170 / 255]),
|
||||
):
|
||||
with self.subTest(hex_code=hex_code):
|
||||
image, = self.node.create_image(3, 2, hex_code)
|
||||
self.assertEqual(image.shape, (1, 2, 3, 3))
|
||||
self.assertEqual(image.dtype, torch.float32)
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
image,
|
||||
torch.tensor(expected, dtype=torch.float32).view(1, 1, 1, 3).expand_as(image),
|
||||
)
|
||||
)
|
||||
|
||||
def test_invalid_hex_colors_raise_a_clear_error(self):
|
||||
for hex_code in ("", "#", "12", "1234", "12345", "1234567", "GG0000"):
|
||||
with self.subTest(hex_code=hex_code):
|
||||
with self.assertRaisesRegex(ValueError, "Invalid HEX color"):
|
||||
self.node.create_image(1, 1, hex_code)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,100 @@
|
||||
import importlib.util
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
MODULE_PATH = Path(__file__).parents[1] / "nodes.py"
|
||||
SPEC = importlib.util.spec_from_file_location("pixelgrid_nodes", MODULE_PATH)
|
||||
NODES = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(NODES)
|
||||
|
||||
|
||||
class GridMedianFixerTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = NODES.GridMedianFixer()
|
||||
|
||||
def test_selects_nonzero_phase_and_removes_single_pixel_noise(self):
|
||||
grid_size = 3
|
||||
colors = torch.tensor(
|
||||
[
|
||||
[
|
||||
[[0.1, 0.2, 0.3], [0.2, 0.4, 0.6], [0.3, 0.6, 0.9]],
|
||||
[[0.4, 0.2, 0.1], [0.5, 0.4, 0.2], [0.6, 0.6, 0.3]],
|
||||
[[0.7, 0.1, 0.4], [0.8, 0.3, 0.5], [0.9, 0.5, 0.6]],
|
||||
]
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
image = torch.ones((1, 12, 12, 3), dtype=torch.float32)
|
||||
for row in range(3):
|
||||
for column in range(3):
|
||||
top = 1 + row * grid_size
|
||||
left = 2 + column * grid_size
|
||||
image[:, top:top + grid_size, left:left + grid_size, :] = colors[
|
||||
:, row:row + 1, column:column + 1, :
|
||||
]
|
||||
|
||||
image[0, 2, 3, :] = 0.0
|
||||
|
||||
downscaled, upscaled = self.node.process_grid(image, grid_size)
|
||||
|
||||
self.assertTrue(torch.equal(downscaled, colors))
|
||||
self.assertEqual(downscaled.shape, (1, 3, 3, 3))
|
||||
self.assertEqual(upscaled.shape, (1, 9, 9, 3))
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
upscaled,
|
||||
colors.repeat_interleave(grid_size, dim=1).repeat_interleave(
|
||||
grid_size, dim=2
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def test_breaks_equal_scores_in_favor_of_top_left_phase(self):
|
||||
image = torch.full((1, 13, 14, 3), 0.4, dtype=torch.float32)
|
||||
|
||||
downscaled, upscaled = self.node.process_grid(image, grid_size=3)
|
||||
|
||||
self.assertEqual(downscaled.shape, (1, 4, 4, 3))
|
||||
self.assertEqual(upscaled.shape, (1, 12, 12, 3))
|
||||
self.assertTrue(torch.all(downscaled == 0.4))
|
||||
self.assertTrue(torch.all(upscaled == 0.4))
|
||||
|
||||
def test_grid_size_one_preserves_image_and_batch(self):
|
||||
image = torch.rand((2, 5, 7, 3), dtype=torch.float32)
|
||||
|
||||
downscaled, upscaled = self.node.process_grid(image, grid_size=1)
|
||||
|
||||
self.assertTrue(torch.equal(downscaled, image))
|
||||
self.assertTrue(torch.equal(upscaled, image))
|
||||
|
||||
def test_image_smaller_than_grid_returns_empty_complete_block_dimensions(self):
|
||||
image = torch.rand((1, 2, 4, 3), dtype=torch.float32)
|
||||
|
||||
downscaled, upscaled = self.node.process_grid(image, grid_size=4)
|
||||
|
||||
self.assertEqual(downscaled.shape, (1, 0, 1, 3))
|
||||
self.assertEqual(upscaled.shape, (1, 0, 4, 3))
|
||||
|
||||
def test_rejects_nonpositive_grid_size(self):
|
||||
image = torch.rand((1, 4, 4, 3), dtype=torch.float32)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "grid_size must be at least 1"):
|
||||
self.node.process_grid(image, grid_size=0)
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
|
||||
def test_keeps_cuda_images_on_device(self):
|
||||
image = torch.rand((1, 8, 8, 3), dtype=torch.float32, device="cuda")
|
||||
|
||||
downscaled, upscaled = self.node.process_grid(image, grid_size=2)
|
||||
|
||||
self.assertEqual(downscaled.device.type, "cuda")
|
||||
self.assertEqual(upscaled.device.type, "cuda")
|
||||
self.assertEqual(downscaled.dtype, image.dtype)
|
||||
self.assertEqual(upscaled.dtype, image.dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user