HakoniwaandClaude Opus 5.5
90e1371b76
- textbbox was measured at the patch position and that position was then subtracted again, so every number was drawn inside patch 0. Measure at the origin and subtract only the glyph offset. - Fall back to DejaVu Sans / Helvetica, then Pillow's scalable default font, so font_size works where Arial is not installed. - Add regression tests for label placement and font size. - Add the Linux/Windows CPU test workflow that the README already describes. - Add .gitignore for __pycache__, bump version to 1.2.1, update README and VALIDATION.md. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01K5PqyHqy9udCb9y3WWmeTx
264 lines
14 KiB
Python
264 lines
14 KiB
Python
import random
|
|
import colorsys
|
|
import json
|
|
from pathlib import Path
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from nodes import ColorshiftColorNode, CsCFill, PaletteEditorNode
|
|
|
|
|
|
torch.set_num_threads(1)
|
|
|
|
|
|
class RegressionTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.palette = torch.tensor([[1., 0., 0.], [0., 0., 1.]])
|
|
self.labels = torch.tensor([[[0, 1], [1, 0]]])
|
|
self.images = self.palette[self.labels]
|
|
|
|
def edit(self, **kwargs):
|
|
options = dict(seed=123, hue_random_enable=False, hue_shift=0.,
|
|
saturation_random_enable=False, saturation_scale=1.,
|
|
value_random_enable=False, value_scale=1.)
|
|
options.update(kwargs)
|
|
return PaletteEditorNode().process_palette(
|
|
self.images, self.palette, self.labels, **options)
|
|
|
|
def test_edit_does_not_mutate_upstream_palette(self):
|
|
original = self.palette.clone()
|
|
self.edit(operations='[{"index": 0, "color": [0, 1, 0]}]')
|
|
torch.testing.assert_close(self.palette, original)
|
|
torch.testing.assert_close(self.edit()[0], self.images)
|
|
|
|
def test_random_edit_repeats_with_seed(self):
|
|
first = self.edit(hue_random_enable=True)[0]
|
|
second = self.edit(hue_random_enable=True)[0]
|
|
torch.testing.assert_close(first, second)
|
|
|
|
def test_random_edit_preserves_global_rng(self):
|
|
state = random.getstate()
|
|
self.edit(hue_random_enable=True)
|
|
self.assertEqual(random.getstate(), state)
|
|
|
|
def test_tiny_image_can_use_large_palette_request(self):
|
|
image = torch.tensor([[[[.2, .4, .6]]]])
|
|
output, palette, labels, _ = ColorshiftColorNode().process(image, 64)
|
|
torch.testing.assert_close(output, image)
|
|
self.assertEqual(palette.shape, (1, 3))
|
|
self.assertEqual(labels.item(), 0)
|
|
|
|
def test_spatial_mask_preserves_original_without_mutation(self):
|
|
image = self.images * .7
|
|
original = image.clone()
|
|
mask = torch.tensor([[[1., 0.], [.5, 0.]]])
|
|
output, _, _, _ = ColorshiftColorNode().process(
|
|
image, 2, lock_masks=mask, palette_override=self.palette)
|
|
expected = image * mask[..., None] + self.images * (1 - mask[..., None])
|
|
torch.testing.assert_close(output, expected)
|
|
torch.testing.assert_close(image, original)
|
|
|
|
def test_fill_honors_explicit_direction_even_when_target_is_darker(self):
|
|
palette = torch.tensor([[.1, .1, .1], [.9, .9, .9]])
|
|
images = palette[self.labels]
|
|
filled, shadow = CsCFill().process(
|
|
images, palette, self.labels, '[{"A": 0, "B": 1}]')
|
|
torch.testing.assert_close(filled, palette[0].expand_as(images))
|
|
torch.testing.assert_close(shadow[..., 3], (self.labels == 1).float())
|
|
|
|
def test_random_edit_changes_with_seed(self):
|
|
self.assertFalse(torch.equal(self.edit(seed=1, hue_random_enable=True)[1],
|
|
self.edit(seed=2, hue_random_enable=True)[1]))
|
|
|
|
def test_mask_inversion_and_explicit_palette_override(self):
|
|
for invert in (False, True):
|
|
with self.subTest(invert=invert):
|
|
result, palette, mask = self.edit(mask_enable=True, lock_color_num="0,",
|
|
invert_mask=invert, hue_shift=120)
|
|
expected = (self.labels == (1 if invert else 0))
|
|
torch.testing.assert_close(mask, expected.float())
|
|
torch.testing.assert_close(result[expected], self.images[expected])
|
|
torch.testing.assert_close(palette[int(invert)], self.palette[int(invert)])
|
|
result, palette, mask = self.edit(mask_enable=True,
|
|
operations='[{"index": 0, "color": [0, 1, 0]}]')
|
|
torch.testing.assert_close(result[self.labels == 0], self.images[self.labels == 0])
|
|
torch.testing.assert_close(palette[0], torch.tensor([0., 1., 0.]))
|
|
|
|
def test_hsv_conversion_matches_standard_library(self):
|
|
colors = torch.cat((torch.rand(100, 3, generator=torch.Generator().manual_seed(9)),
|
|
torch.tensor([[0., 0., 0.], [1., 1., 1.], [.5, .5, .5]])))
|
|
editor = PaletteEditorNode()
|
|
hsv = editor.rgb_to_hsv(colors)
|
|
expected = torch.tensor([colorsys.rgb_to_hsv(*color.tolist()) for color in colors])
|
|
torch.testing.assert_close(hsv, expected)
|
|
torch.testing.assert_close(editor.hsv_to_rgb(hsv), colors)
|
|
|
|
def test_hsv_operations_and_rgb_precedence(self):
|
|
_, palette, _ = self.edit(operations='[{"index": 0, "color": [1, 0, 0], "hsv": [0.5, 1, 1]}]')
|
|
torch.testing.assert_close(palette[0], torch.tensor([0., 1., 1.]))
|
|
|
|
def test_invalid_operations_raise_without_mutation(self):
|
|
invalid = ['oops', '{}', '[null]', '[{}]', '[{"index": -1, "color": [0,0,0]}]',
|
|
'[{"index": 0.5, "color": [0,0,0]}]', '[{"index": true, "color": [0,0,0]}]',
|
|
'[{"index": 0, "color": [2,0,0]}]', '[{"index": 0, "color": [NaN,0,0]}]',
|
|
'[{"index": 0, "color": [0,0]}]', '[{"index": 0, "color": ["0",0,0]}]',
|
|
'[{"index": 0, "hsv": null}]', '[{"index": 0}]',
|
|
'[{"index": 0, "color": [0,1,0]}, {"index": 999, "color": [0,0,0]}]']
|
|
for operation in invalid:
|
|
with self.subTest(operation=operation), self.assertRaisesRegex(ValueError, 'operations|palette operation'):
|
|
self.edit(operations=operation)
|
|
torch.testing.assert_close(self.images, self.palette[self.labels])
|
|
|
|
def test_invalid_locks_are_actionable(self):
|
|
for locks in ('x', '-1', '999'):
|
|
with self.subTest(locks=locks), self.assertRaisesRegex(ValueError, 'lock_color_num'):
|
|
self.edit(mask_enable=True, lock_color_num=locks)
|
|
|
|
def test_bundled_editor_workflows_and_comments(self):
|
|
for path in (Path(__file__).resolve().parents[1] / 'example').glob('*.json'):
|
|
for node in json.loads(path.read_text(encoding='utf8'))['nodes']:
|
|
if node['type'] == 'CsCPaletteEditor':
|
|
with self.subTest(workflow=path.name):
|
|
# The supplied commented example must behave exactly like [].
|
|
torch.testing.assert_close(self.edit(operations=node['widgets_values'][-1])[0], self.images)
|
|
torch.testing.assert_close(self.edit(operations=' ')[0], self.images)
|
|
|
|
def test_quantization_seed_and_global_numpy_rng(self):
|
|
image = torch.rand(2, 32, 32, 3, generator=torch.Generator().manual_seed(10))
|
|
state = np.random.get_state()
|
|
first = ColorshiftColorNode().process(image, 4, seed=2**64 - 1)
|
|
second = ColorshiftColorNode().process(image, 4, seed=2**64 - 1)
|
|
for a, b in zip(first, second):
|
|
torch.testing.assert_close(a, b)
|
|
after = np.random.get_state()
|
|
self.assertEqual(state[0], after[0])
|
|
np.testing.assert_array_equal(state[1], after[1])
|
|
self.assertEqual(state[2:], after[2:])
|
|
|
|
def test_solid_and_thin_images(self):
|
|
for shape in ((1, 1, 20, 3), (1, 20, 1, 3), (2, 4, 4, 3)):
|
|
with self.subTest(shape=shape):
|
|
image = torch.full(shape, .25)
|
|
output, palette, labels, _ = ColorshiftColorNode().process(image, 64, sampling_rate=.01)
|
|
torch.testing.assert_close(output, image)
|
|
self.assertEqual(len(palette), 1)
|
|
self.assertEqual(labels.max().item(), 0)
|
|
|
|
def test_palette_is_sorted_by_coverage(self):
|
|
image = torch.tensor([[[[1., 0., 0.]] * 3 + [[0., 0., 1.]]]])
|
|
output, palette, labels, _ = ColorshiftColorNode().process(image, 64)
|
|
torch.testing.assert_close(output, image)
|
|
torch.testing.assert_close(palette[0], torch.tensor([1., 0., 0.]))
|
|
self.assertEqual(labels.tolist(), [[[0, 0, 0, 1]]])
|
|
|
|
def test_palette_override_order_and_float16_images(self):
|
|
image = self.images.half().transpose(1, 2)
|
|
result, palette, labels, _ = ColorshiftColorNode().process(image, 8, palette_override=self.palette)
|
|
torch.testing.assert_close(result, image)
|
|
torch.testing.assert_close(palette, self.palette)
|
|
self.assertEqual(labels.dtype, torch.long)
|
|
|
|
def test_palette_matching_is_chunked_and_exact(self):
|
|
pixels = torch.rand(2, 40000, 3, generator=torch.Generator().manual_seed(12))
|
|
original_cdist = torch.cdist
|
|
with patch('nodes.torch.cdist', wraps=original_cdist) as calls:
|
|
labels = ColorshiftColorNode()._match_palette_batch(pixels, self.palette)
|
|
self.assertGreater(calls.call_count, 1)
|
|
self.assertLessEqual(max(call.args[0].shape[0] for call in calls.call_args_list), 65536)
|
|
distances = ((pixels.unsqueeze(-2) - self.palette) ** 2).sum(-1)
|
|
torch.testing.assert_close(labels, distances.argmin(-1))
|
|
|
|
def test_mask_broadcast_resize_and_repeat_last(self):
|
|
images = (self.images * .7).repeat(3, 1, 1, 1)
|
|
masks = torch.tensor([[[1.]], [[0.]]])
|
|
output = ColorshiftColorNode().process(images, 2, lock_masks=masks, palette_override=self.palette)[0]
|
|
torch.testing.assert_close(output[0], images[0])
|
|
torch.testing.assert_close(output[1:], self.images.expand(2, -1, -1, -1))
|
|
output = ColorshiftColorNode().process(images, 2, lock_masks=torch.ones(2, 2), palette_override=self.palette)[0]
|
|
torch.testing.assert_close(output, images)
|
|
|
|
def test_invalid_spatial_masks(self):
|
|
for mask in (torch.empty(0, 2, 2), torch.ones(2), torch.full((2, 2), float('nan'))):
|
|
with self.subTest(shape=mask.shape), self.assertRaisesRegex(ValueError, 'lock_masks'):
|
|
ColorshiftColorNode().process(self.images, 2, palette_override=self.palette, lock_masks=mask)
|
|
|
|
def test_invalid_images_and_palettes(self):
|
|
for image in (torch.ones(1, 2, 2, 4), torch.empty(0, 2, 2, 3), self.images.long(), self.images * float('nan')):
|
|
with self.subTest(shape=image.shape), self.assertRaisesRegex(ValueError, 'images'):
|
|
ColorshiftColorNode().process(image, 2)
|
|
for palette in (torch.empty(0, 3), torch.ones(3), self.palette * float('inf'), self.palette * 2):
|
|
with self.subTest(palette=palette), self.assertRaisesRegex(ValueError, 'palette'):
|
|
ColorshiftColorNode().process(self.images, 2, palette_override=palette)
|
|
|
|
def test_invalid_index_maps_in_both_consumers(self):
|
|
for labels in (self.labels - 1, self.labels + 2, self.labels.float() + .5,
|
|
self.labels.float() * float('nan'), self.labels[..., :1]):
|
|
for process in (PaletteEditorNode().process_palette, CsCFill().process):
|
|
with self.subTest(labels=labels, process=process), self.assertRaisesRegex(ValueError, 'index_maps'):
|
|
process(self.images, self.palette, labels)
|
|
|
|
def test_fill_rejects_invalid_manual_pairs_instead_of_auto_filling(self):
|
|
for operation in ('bad json', '{}', '[{}]', '[{"A": 0, "B": 999}]', '[{"A": 0.5, "B": 1}]'):
|
|
with self.subTest(operation=operation), self.assertRaisesRegex(ValueError, 'operations'):
|
|
CsCFill().process(self.images, self.palette, self.labels, operation)
|
|
|
|
def test_fill_shadow_rgb_and_transparent_background(self):
|
|
_, shadow = CsCFill().process(self.images, self.palette, self.labels, '[{"A": 0, "B": 1}]')
|
|
selected = self.labels == 1
|
|
torch.testing.assert_close(shadow[..., :3][selected], self.images[selected])
|
|
self.assertEqual(torch.count_nonzero(shadow[~selected]).item(), 0)
|
|
|
|
def test_fill_auto_pairs_odd_palette_and_self_pair(self):
|
|
palette = torch.tensor([[.1, .1, .1], [.2, .2, .2], [1., 0., 0.]])
|
|
self.assertEqual(CsCFill()._compute_auto_pairs(palette), [(1, 0)])
|
|
output, shadow = CsCFill().process(self.images, self.palette, self.labels,
|
|
'[{"A": 0, "B": 1}, {"A": 1, "B": 1}]')
|
|
torch.testing.assert_close(output, self.images)
|
|
self.assertEqual(torch.count_nonzero(shadow).item(), 0)
|
|
|
|
def test_preview_has_readable_text_on_white(self):
|
|
preview = ColorshiftColorNode().generate_palette_preview(torch.ones(1, 3), 20)
|
|
self.assertEqual(preview.shape, (1, 100, 800, 3))
|
|
self.assertLess(preview[0, 25:75, 25:75].min().item(), .5)
|
|
|
|
def test_preview_labels_every_patch_in_place(self):
|
|
# Every patch must carry its own number; previously all labels stacked on patch 0.
|
|
preview = ColorshiftColorNode().generate_palette_preview(torch.ones(12, 3), 30)
|
|
self.assertEqual(preview.shape, (1, 200, 800, 3))
|
|
for i in range(12):
|
|
x, y = (i % 8) * 100, (i // 8) * 100
|
|
with self.subTest(patch=i):
|
|
self.assertLess(preview[0, y + 20:y + 80, x + 20:x + 80].min().item(), .5)
|
|
self.assertEqual(preview[0, y:y + 10, x:x + 100].min().item(), 1.)
|
|
|
|
def test_preview_font_size_is_honored(self):
|
|
def ink(size):
|
|
preview = ColorshiftColorNode().generate_palette_preview(torch.ones(1, 3), size)
|
|
return (preview[0, :, :100] < .5).any(-1).sum().item()
|
|
self.assertGreater(ink(50), ink(10) * 4)
|
|
|
|
@unittest.skipUnless(torch.cuda.is_available(), 'CUDA is not available')
|
|
def test_all_nodes_on_cuda_with_cpu_auxiliary_inputs(self):
|
|
for dtype in (torch.float32, torch.float16, torch.bfloat16):
|
|
with self.subTest(dtype=dtype):
|
|
images = self.images.to(device='cuda', dtype=dtype)
|
|
output, palette, labels, _ = ColorshiftColorNode().process(
|
|
images, 2, palette_override=self.palette, lock_masks=torch.ones(2, 2))
|
|
torch.testing.assert_close(output, images)
|
|
result, edited, mask = PaletteEditorNode().process_palette(
|
|
images, self.palette, self.labels, operations='[{"index": 0, "color": [0, 1, 0]}]')
|
|
filled, shadow = CsCFill().process(images, self.palette, self.labels, '[{"A": 0, "B": 1}]')
|
|
for value in (output, palette, labels, result, edited, mask, filled, shadow):
|
|
self.assertEqual(value.device.type, 'cuda')
|
|
self.assertEqual(result.dtype, dtype)
|
|
self.assertEqual(shadow.dtype, dtype)
|
|
generated = ColorshiftColorNode().process(self.images.cuda(), 64)
|
|
torch.testing.assert_close(generated[0], self.images.cuda())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|