ausboss
f2a512acd7
_fill_tensor already returns the pad color sized to the image's channels, with extra channels opaque, so _color_pad's own alpha branch could never run; a test now pins the opaque alpha instead. Save Image confines every destination to the output folder before writing, so the preview loop's "saved outside the output folder" fallback was dead too.
160 lines
7.0 KiB
Python
160 lines
7.0 KiB
Python
"""Offline tests for Align Image and Image Size."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes._align_helpers import ALIGN_MODES, align_image, aligned_size
|
|
from nodes.node_align_image import AusBossAlignImage
|
|
from nodes.node_image_size import AusBossImageSize
|
|
|
|
|
|
class AlignedSizeTests(unittest.TestCase):
|
|
def test_crop_rounds_down_pad_rounds_up_resize_rounds_nearest(self):
|
|
# 1000 / 16 = 62.5: down 992, up 1008, nearest ties-up 1008.
|
|
self.assertEqual(aligned_size(1000, 1000, 16, "crop"), (992, 992))
|
|
self.assertEqual(aligned_size(1000, 1000, 16, "pad"), (1008, 1008))
|
|
self.assertEqual(aligned_size(1000, 1000, 16, "resize"), (1008, 1008))
|
|
# 1030 is nearer 1024 than 1040.
|
|
self.assertEqual(aligned_size(1030, 500, 16, "resize"), (1024, 496))
|
|
|
|
def test_already_aligned_sizes_pass_through_every_mode(self):
|
|
for mode in ALIGN_MODES:
|
|
self.assertEqual(aligned_size(1024, 768, 32, mode), (1024, 768))
|
|
|
|
def test_a_side_smaller_than_one_multiple_snaps_up_to_one(self):
|
|
# Nothing can crop below the smallest legal size.
|
|
for mode in ALIGN_MODES:
|
|
self.assertEqual(aligned_size(20, 100, 32, mode)[0], 32, mode)
|
|
|
|
def test_multiple_one_changes_nothing(self):
|
|
for mode in ALIGN_MODES:
|
|
self.assertEqual(aligned_size(1013, 777, 1, mode), (1013, 777))
|
|
|
|
def test_unknown_mode_raises(self):
|
|
with self.assertRaises(ValueError):
|
|
aligned_size(100, 100, 16, "stretch")
|
|
|
|
|
|
def ramp(batch: int, height: int, width: int) -> torch.Tensor:
|
|
"""A BHWC batch whose pixel values encode their own x coordinate."""
|
|
xs = torch.arange(width, dtype=torch.float32) / max(1, width - 1)
|
|
return xs.view(1, 1, -1, 1).expand(batch, height, width, 3).clone()
|
|
|
|
|
|
class AlignImageTests(unittest.TestCase):
|
|
def test_resize_hits_the_nearest_multiple(self):
|
|
out, width, height, ox, oy = align_image(ramp(2, 500, 1000), 16, "resize")
|
|
self.assertEqual((width, height), (1008, 496))
|
|
self.assertEqual(tuple(out.shape), (2, 496, 1008, 3))
|
|
self.assertEqual((ox, oy), (0, 0))
|
|
|
|
def test_crop_is_centered_and_lossless_in_the_kept_region(self):
|
|
frames = ramp(1, 100, 100)
|
|
out, width, height, ox, oy = align_image(frames, 32, "crop")
|
|
self.assertEqual((width, height), (96, 96))
|
|
# Center crop of 100 -> 96 drops 2 columns each side.
|
|
self.assertTrue(torch.equal(out, frames[:, 2:98, 2:98, :]))
|
|
# The original's corner now sits 2 px outside the output.
|
|
self.assertEqual((ox, oy), (-2, -2))
|
|
|
|
def test_pad_replicates_edges_up_to_the_next_multiple(self):
|
|
frames = ramp(1, 100, 100)
|
|
out, width, height, ox, oy = align_image(frames, 32, "pad")
|
|
self.assertEqual((width, height), (128, 128))
|
|
# 28 extra pixels split 14/14; the original sits centered, untouched.
|
|
self.assertTrue(torch.equal(out[:, 14:114, 14:114, :], frames))
|
|
# The offsets point exactly at where the original landed.
|
|
self.assertEqual((ox, oy), (14, 14))
|
|
self.assertTrue(
|
|
torch.equal(out[:, oy : oy + 100, ox : ox + 100, :], frames)
|
|
)
|
|
# Padding replicates the border, so corners equal the corner pixel.
|
|
self.assertTrue(torch.equal(out[0, 0, 0], frames[0, 0, 0]))
|
|
self.assertTrue(torch.equal(out[0, -1, -1], frames[0, -1, -1]))
|
|
|
|
def test_already_aligned_input_is_returned_as_the_same_object(self):
|
|
frames = ramp(1, 64, 128)
|
|
out, width, height, _ox, _oy = align_image(frames, 32, "crop")
|
|
self.assertIs(out, frames)
|
|
self.assertEqual((width, height), (128, 64))
|
|
|
|
def test_crop_mode_pads_a_side_it_cannot_crop(self):
|
|
# 20 px wide with multiple 32: crop cannot reach a multiple below,
|
|
# so that side replicate-pads up to exactly one multiple.
|
|
out, width, height, ox, oy = align_image(ramp(1, 100, 20), 32, "crop")
|
|
self.assertEqual((width, height), (32, 96))
|
|
self.assertEqual(tuple(out.shape), (1, 96, 32, 3))
|
|
# Mixed axes: the padded side is positive, the cropped side negative.
|
|
self.assertEqual((ox, oy), (6, -2))
|
|
|
|
def test_color_pad_keeps_an_alpha_channel_opaque(self):
|
|
frames = torch.cat([ramp(1, 20, 20), torch.full((1, 20, 20, 1), 0.25)], dim=-1)
|
|
out, width, height, ox, oy = align_image(
|
|
frames, 16, "pad", pad_fill="color", pad_color="#ff0000"
|
|
)
|
|
self.assertEqual((width, height), (32, 32))
|
|
self.assertEqual(tuple(out.shape), (1, 32, 32, 4))
|
|
self.assertTrue(torch.equal(out[:, oy : oy + 20, ox : ox + 20, :], frames))
|
|
# The padded border takes the color, and the alpha channel pads opaque.
|
|
self.assertTrue(torch.equal(out[0, 0, 0], torch.tensor([1.0, 0.0, 0.0, 1.0])))
|
|
self.assertTrue(torch.equal(out[0, -1, -1], torch.tensor([1.0, 0.0, 0.0, 1.0])))
|
|
|
|
def test_batch_and_values_survive(self):
|
|
frames = ramp(3, 40, 40)
|
|
out, _w, _h, _ox, _oy = align_image(frames, 16, "pad")
|
|
self.assertEqual(out.shape[0], 3)
|
|
self.assertGreaterEqual(float(out.min()), 0.0)
|
|
self.assertLessEqual(float(out.max()), 1.0)
|
|
|
|
def test_rejects_non_bhwc_input(self):
|
|
with self.assertRaises(ValueError):
|
|
align_image(torch.zeros((64, 64, 3)), 16, "resize")
|
|
|
|
|
|
class AlignImageNodeTests(unittest.TestCase):
|
|
def test_node_returns_image_and_size_only(self):
|
|
# The helper still reports offsets; the node's face stops at the
|
|
# three outputs a graph wires (offsets were dropped 2026-09-03).
|
|
node = AusBossAlignImage()
|
|
result = node.align(ramp(1, 100, 100), 32, "pad")
|
|
self.assertEqual(len(result), 3)
|
|
out, width, height = result
|
|
self.assertEqual((width, height), (128, 128))
|
|
self.assertEqual(tuple(out.shape), (1, 128, 128, 3))
|
|
for value in (width, height):
|
|
self.assertIsInstance(value, int)
|
|
self.assertEqual(AusBossAlignImage.RETURN_NAMES, ("image", "width", "height"))
|
|
|
|
def test_mode_choices_match_the_helper(self):
|
|
widget = AusBossAlignImage.INPUT_TYPES()["required"]["mode"][0]
|
|
self.assertEqual(tuple(widget), ALIGN_MODES)
|
|
|
|
|
|
class ImageSizeNodeTests(unittest.TestCase):
|
|
def test_reports_width_height_and_both_edges(self):
|
|
node = AusBossImageSize()
|
|
self.assertEqual(node.measure(torch.zeros((2, 480, 832, 3))), (832, 480, 832, 480, 2))
|
|
self.assertEqual(node.measure(torch.zeros((1, 1080, 608, 3))), (608, 1080, 1080, 608, 1))
|
|
|
|
def test_square_images_agree_on_both_edges(self):
|
|
node = AusBossImageSize()
|
|
self.assertEqual(node.measure(torch.zeros((1, 512, 512, 3))), (512, 512, 512, 512, 1))
|
|
|
|
def test_rejects_non_bhwc_input(self):
|
|
with self.assertRaises(ValueError):
|
|
AusBossImageSize().measure(torch.zeros((512, 512)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|