From 005997b1ae8e666e5068866c4e97fd72a1b486b8 Mon Sep 17 00:00:00 2001 From: molbal Date: Mon, 5 Oct 2026 22:20:01 +0200 Subject: [PATCH] Improve median fixer algo --- __init__.py | 5 +- nodes.py | 171 ++++++++++++++++++++++++++++--- pyproject.toml | 2 +- readme.md | 33 +++++- tests/test_create_solid_color.py | 52 ++++++++++ tests/test_grid_median_fixer.py | 100 ++++++++++++++++++ 6 files changed, 345 insertions(+), 18 deletions(-) create mode 100644 tests/test_create_solid_color.py create mode 100644 tests/test_grid_median_fixer.py diff --git a/__init__.py b/__init__.py index bd351ef..3c9fc24 100644 --- a/__init__.py +++ b/__init__.py @@ -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" } diff --git a/nodes.py b/nodes.py index 21e4fcc..d94ac8e 100644 --- a/nodes.py +++ b/nodes.py @@ -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) \ No newline at end of file + 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) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index ed1a203..61e8ef8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = [] diff --git a/readme.md b/readme.md index e3665b6..aa7aeb4 100644 --- a/readme.md +++ b/readme.md @@ -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 diff --git a/tests/test_create_solid_color.py b/tests/test_create_solid_color.py new file mode 100644 index 0000000..3dfc673 --- /dev/null +++ b/tests/test_create_solid_color.py @@ -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() diff --git a/tests/test_grid_median_fixer.py b/tests/test_grid_median_fixer.py new file mode 100644 index 0000000..70ba46d --- /dev/null +++ b/tests/test_grid_median_fixer.py @@ -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()