ausboss
8837c75553
Stitch Inpaint takes an RGB patch over an RGBA picture and keeps the picture's own alpha, instead of stopping on a channel mismatch. The halo fix spreads colours only. Image Compare keeps alpha in a preview that has see-through parts and flags it, so the panel shows them over a checkerboard (Slide, A/B and Erase). The Erase mix lays the top picture over the other, the saved PNG is RGBA, and the PSD carries each layer's transparency plus a merged alpha, its flat colours stored over white as Photoshop does. A see-through B is resized premultiplied. Align measures over gray and warps all four channels. Solid pictures give byte-identical files. Refs #232
122 lines
5.7 KiB
Python
122 lines
5.7 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
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 import _preview_helpers, node_compare
|
|
from nodes._compare_helpers import pair_pictures
|
|
from nodes.node_compare import AusBossCompare, first_frame_to_pil
|
|
|
|
|
|
def fake_folder_paths(temp_dir: str):
|
|
class FakeFolderPaths:
|
|
@staticmethod
|
|
def get_temp_directory():
|
|
return temp_dir
|
|
|
|
@staticmethod
|
|
def get_save_image_path(prefix, output_dir, _width, _height):
|
|
folder = Path(output_dir) / "ausboss"
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
return str(folder), prefix, 1, "ausboss", prefix
|
|
|
|
return FakeFolderPaths
|
|
|
|
|
|
class CompareNodeTests(unittest.TestCase):
|
|
def test_first_frame_conversion_and_validation(self):
|
|
batch = torch.zeros(3, 4, 6, 3)
|
|
batch[0, :, :, 0] = 1.0 # first frame is pure red
|
|
preview = first_frame_to_pil(batch)
|
|
self.assertEqual((preview.width, preview.height), (6, 4))
|
|
self.assertEqual(preview.getpixel((0, 0)), (255, 0, 0))
|
|
with self.assertRaisesRegex(ValueError, "BHWC"):
|
|
first_frame_to_pil(torch.zeros(4, 6, 3))
|
|
with self.assertRaisesRegex(ValueError, "empty"):
|
|
first_frame_to_pil(torch.zeros(0, 4, 6, 3))
|
|
|
|
def test_compare_saves_temp_previews_and_passes_both_through(self):
|
|
image_a = torch.rand(2, 8, 12, 3)
|
|
image_b = torch.rand(1, 8, 12, 3)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
# folder_paths is imported by the shared preview helper now, which
|
|
# is where every node's temp preview is written.
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)), \
|
|
patch.object(node_compare, "folder_paths", fake_folder_paths(tmp)):
|
|
result = AusBossCompare().compare(image_a, image_b)
|
|
|
|
# image_a stays the first output, so saved links keep their slot;
|
|
# image_b was appended after it.
|
|
self.assertEqual(AusBossCompare.RETURN_NAMES, ("image_a", "image_b"))
|
|
self.assertEqual(len(AusBossCompare.OUTPUT_TOOLTIPS), 2)
|
|
self.assertIs(result["result"][0], image_a)
|
|
self.assertIs(result["result"][1], image_b)
|
|
for key in ("a_images", "b_images"):
|
|
(ref,) = result["ui"][key]
|
|
self.assertEqual(ref["type"], "temp")
|
|
self.assertEqual(ref["subfolder"], "ausboss")
|
|
self.assertEqual((ref["width"], ref["height"]), (12, 8))
|
|
saved = Path(tmp) / "ausboss" / ref["filename"]
|
|
with Image.open(saved) as png:
|
|
self.assertEqual(png.size, (12, 8))
|
|
# The run's key stands for exactly these two files; Erase's Save
|
|
# sends the key back and never a file name.
|
|
(key,) = result["ui"]["pair"]
|
|
self.assertEqual(
|
|
pair_pictures(key)[:2],
|
|
tuple(Path(tmp) / "ausboss" / result["ui"][name][0]["filename"] for name in ("a_images", "b_images")),
|
|
)
|
|
# No name wired in: the one typed on the node will be used.
|
|
self.assertEqual((result["ui"]["save_name"], pair_pictures(key)[2]), ([""], ""))
|
|
|
|
def test_a_see_through_picture_keeps_its_clear_parts_in_the_preview(self):
|
|
image_a = torch.ones(1, 8, 12, 4)
|
|
image_a[0, :, :4, 3] = 0.0
|
|
image_b = torch.rand(1, 8, 12, 3)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)), \
|
|
patch.object(node_compare, "folder_paths", fake_folder_paths(tmp)):
|
|
result = AusBossCompare().compare(image_a, image_a)
|
|
(ref_a,) = result["ui"]["a_images"]
|
|
with Image.open(Path(tmp) / "ausboss" / ref_a["filename"]) as png:
|
|
self.assertEqual(png.mode, "RGBA")
|
|
self.assertEqual(png.getpixel((0, 0))[3], 0)
|
|
result = AusBossCompare().compare(image_a, image_b)
|
|
# The panel shows the clear parts over a checkerboard; an opaque B has no flag.
|
|
self.assertTrue(result["ui"]["a_images"][0]["alpha"])
|
|
self.assertNotIn("alpha", result["ui"]["b_images"][0])
|
|
|
|
def test_a_wired_filename_names_the_saved_files(self):
|
|
image = torch.rand(1, 8, 12, 3)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with patch.object(_preview_helpers, "folder_paths", fake_folder_paths(tmp)), \
|
|
patch.object(node_compare, "folder_paths", fake_folder_paths(tmp)):
|
|
result = AusBossCompare().compare(image, image, filename="sets/a/photo 12.jpg")
|
|
self.assertEqual(result["ui"]["save_name"], ["photo 12"])
|
|
self.assertEqual(pair_pictures(result["ui"]["pair"][0])[2], "photo 12")
|
|
self.assertIs(result["result"][0], image)
|
|
optional = AusBossCompare.INPUT_TYPES()["optional"]["filename"]
|
|
self.assertEqual(optional[0], "STRING")
|
|
self.assertTrue(optional[1]["forceInput"])
|
|
self.assertEqual(list(AusBossCompare.INPUT_TYPES()["required"]), ["image_a", "image_b"])
|
|
|
|
def test_two_nodes_use_distinct_temp_prefixes(self):
|
|
# Prefixes are randomized per instance so parallel Compare nodes
|
|
# never overwrite each other's previews.
|
|
self.assertNotEqual(AusBossCompare()._prefix, AusBossCompare()._prefix)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|