Improve median fixer algo

This commit is contained in:
molbal
2026-10-05 22:20:01 +02:00
parent 4c63500aaa
commit 005997b1ae
6 changed files with 345 additions and 18 deletions
+4 -1
View File
@@ -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"
}
+158 -13
View File
@@ -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
View File
@@ -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 = []
+30 -3
View File
@@ -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
+52
View File
@@ -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()
+100
View File
@@ -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()