ausboss
02a07b84f3
A part painted in the MaskEditor (see-through in the saved file) was always handed on as flat fill. The new optional painted_area input, "fill" by default, can be set to "keep picture": the mask is unchanged and the image, original and prompt_image outputs carry the stored picture there. The editor's Inpaint & Stitch section has the switch.
471 lines
25 KiB
Python
471 lines
25 KiB
Python
"""See-through parts of a picture are area to paint.
|
|
|
|
A cutout's see-through pixels store a colour nobody sees - black, or the old
|
|
background - and Image Crop + Rotate + Pad used to blend it over the fill:
|
|
an oval PNG reached the model with a dark ring round it, and Stitch Inpaint
|
|
pasted part of the ring back. Now a pixel less than 90% as solid as the
|
|
picture's most solid pixel is empty, like a corner a turn leaves: fill
|
|
colour in the canvas, white in the prompt view and in `original`, marked in
|
|
the mask, never pasted back. Every kept pixel is fully solid in its own
|
|
colour. A picture at least 90% solid everywhere goes through exactly as
|
|
before; OpaqueIdentityTests pins that against the transform as it was.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
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._inpaint_crop_helpers import ( # noqa: E402
|
|
SEAM_BLEND_IN,
|
|
apply_stitch,
|
|
build_transform_stitcher,
|
|
seam_plan,
|
|
)
|
|
from nodes._transform_engine import ( # noqa: E402
|
|
TransformSpec,
|
|
_geometry,
|
|
_validate_source,
|
|
fill_rgb,
|
|
see_through_kept,
|
|
transform_pil,
|
|
transform_tensor_batch,
|
|
)
|
|
from nodes.node_image_crop_rotate_pad import AusBossImageCropRotatePad # noqa: E402
|
|
from nodes.node_video_crop_rotate_pad import AusBossVideoCropRotatePad # noqa: E402
|
|
|
|
COMFY_ROOT = os.environ.get("AUSBOSS_COMFY_ROOT")
|
|
if COMFY_ROOT:
|
|
sys.path.append(COMFY_ROOT)
|
|
sys.argv = sys.argv[:1] # ComfyUI's argument parser reads sys.argv
|
|
import comfy.cli_args # noqa: E402
|
|
|
|
# blend in checks for a cancel through comfy.model_management, which
|
|
# would otherwise claim the GPU; these checks need none.
|
|
comfy.cli_args.args.cpu = True
|
|
|
|
|
|
def reference_transform(image: Image.Image, spec: TransformSpec):
|
|
"""The transform as it was before see-through parts were painted: the
|
|
source's alpha blended straight over the fill and the mask its inverse.
|
|
Kept verbatim so a picture with no see-through part is held to it."""
|
|
_validate_source(image)
|
|
spec = spec.normalized()
|
|
rgba = image.convert("RGBA")
|
|
if spec.rotation_degrees == 0.0:
|
|
rotated = rgba.copy()
|
|
else:
|
|
rotated = rgba.rotate(
|
|
-spec.rotation_degrees,
|
|
resample=Image.Resampling.BICUBIC,
|
|
expand=True,
|
|
fillcolor=(*fill_rgb(spec.fill_color), 0),
|
|
)
|
|
geometry = _geometry(rotated, spec)
|
|
cropped = rotated.crop((
|
|
geometry.crop_x,
|
|
geometry.crop_y,
|
|
geometry.crop_x + geometry.crop_width,
|
|
geometry.crop_y + geometry.crop_height,
|
|
))
|
|
alpha = cropped.getchannel("A")
|
|
fill = fill_rgb(spec.fill_color)
|
|
filled_crop = Image.new("RGB", cropped.size, fill)
|
|
filled_crop.paste(cropped.convert("RGB"), mask=alpha)
|
|
output = Image.new("RGB", (geometry.output_width, geometry.output_height), fill)
|
|
output.paste(filled_crop, (geometry.pad_left, geometry.pad_top))
|
|
mask = Image.new("L", output.size, 255)
|
|
mask.paste(Image.fromarray(255 - np.asarray(alpha, dtype=np.uint8)), (geometry.pad_left, geometry.pad_top))
|
|
if spec.feather > 0:
|
|
original = np.asarray(mask, dtype=np.uint8)
|
|
blurred = np.asarray(mask.filter(ImageFilter.GaussianBlur(spec.feather)), dtype=np.uint16)
|
|
mask = Image.fromarray(np.maximum(original, np.minimum(blurred * 2, 255).astype(np.uint8)))
|
|
return output, mask, geometry
|
|
|
|
|
|
def as_tensor(image: Image.Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.asarray(image, dtype=np.float32).copy() / 255.0).unsqueeze(0)
|
|
|
|
|
|
def run_node(source: Image.Image, *frames: Image.Image, **values):
|
|
frames = frames or (source,)
|
|
with patch("nodes.node_image_crop_rotate_pad.resolve_input_path", return_value=Path("source.png")), patch(
|
|
"nodes.node_image_crop_rotate_pad.load_image_frames", return_value=[frame.copy() for frame in frames]
|
|
):
|
|
return AusBossImageCropRotatePad().load_transform("source.png", **values)
|
|
|
|
|
|
LAYOUTS = {
|
|
"as loaded": {},
|
|
"padded": dict(pad_left=40, pad_top=12, pad_right=24, pad_bottom=6),
|
|
"padded and feathered": dict(pad_left=32, pad_right=32, feather=24),
|
|
"turned": dict(rotation_degrees=17.3),
|
|
"turned, padded, feathered": dict(rotation_degrees=-9, pad_left=24, pad_right=40, feather=12),
|
|
"turned, cropped and padded": dict(
|
|
rotation_degrees=-18.1, crop_x=10, crop_y=6, crop_width=60, crop_height=50,
|
|
pad_left=30, pad_top=20, pad_right=30, pad_bottom=20, feather=8,
|
|
),
|
|
"ratio and multiple": dict(crop_aspect_ratio="16:9", canvas_multiple=32, pad_left=10),
|
|
"black fill, turned 90": dict(rotation_degrees=90, pad_top=12, fill_color="#000000", feather=6),
|
|
"stitch settings": dict(pad_bottom=30, stitch_blend=8, stitch_grow=-3, feather=10),
|
|
}
|
|
|
|
|
|
def opaque_sources():
|
|
rng = np.random.default_rng(5)
|
|
colours = rng.integers(0, 256, (72, 96, 3), dtype=np.uint8)
|
|
yield "RGB photo", Image.fromarray(colours)
|
|
yield "RGBA, alpha 255", Image.fromarray(np.dstack([colours, np.full((72, 96), 255, np.uint8)]), "RGBA")
|
|
# Generated pictures often carry alpha 249-254 in places; at least 90%
|
|
# solid everywhere means nothing is see-through and nothing changes.
|
|
noise = rng.integers(249, 256, (72, 96), dtype=np.uint8)
|
|
yield "RGBA, alpha 249-255", Image.fromarray(np.dstack([colours, noise]), "RGBA")
|
|
edge = np.full((72, 96), 255, np.uint8)
|
|
edge[:, :3] = 230 # exactly 90% solid: still kept
|
|
yield "RGBA, edge at 90%", Image.fromarray(np.dstack([colours, edge]), "RGBA")
|
|
|
|
|
|
class OpaqueIdentityTests(unittest.TestCase):
|
|
"""A picture with no see-through part comes out exactly as before."""
|
|
|
|
def test_the_transform_matches_the_old_one_bit_for_bit(self):
|
|
for source_name, source in opaque_sources():
|
|
for layout_name, layout in LAYOUTS.items():
|
|
spec = TransformSpec(**{k: v for k, v in layout.items() if not k.startswith("stitch")})
|
|
with self.subTest(source=source_name, layout=layout_name):
|
|
output, mask, geometry = transform_pil(source, spec)
|
|
old_output, old_mask, old_geometry = reference_transform(source, spec)
|
|
self.assertEqual(geometry, old_geometry)
|
|
self.assertEqual(output.tobytes(), old_output.tobytes())
|
|
self.assertEqual(mask.tobytes(), old_mask.tobytes())
|
|
|
|
def test_every_node_output_matches_the_old_one(self):
|
|
model = torch.from_numpy(np.random.default_rng(9).random((1, 200, 200, 3), dtype=np.float32))
|
|
for source_name, source in opaque_sources():
|
|
for layout_name, layout in LAYOUTS.items():
|
|
with self.subTest(source=source_name, layout=layout_name):
|
|
image, mask, stitcher, original, width, height, prompt_image = run_node(source, **layout)
|
|
spec = TransformSpec(**{k: v for k, v in layout.items() if not k.startswith("stitch")})
|
|
old_output, old_mask, geometry = reference_transform(source, spec)
|
|
old_image, old_mask = as_tensor(old_output), as_tensor(old_mask)
|
|
old_stitcher = build_transform_stitcher(
|
|
old_image, old_mask, geometry, layout.get("stitch_blend", 32), layout.get("stitch_grow", 0)
|
|
)
|
|
self.assertTrue(torch.equal(image, old_image))
|
|
self.assertTrue(torch.equal(mask, old_mask))
|
|
self.assertEqual((width, height), (old_output.width, old_output.height))
|
|
for key in ("canvas", "blend", "generated"):
|
|
self.assertTrue(torch.equal(stitcher[key], old_stitcher[key]), key)
|
|
for key in ("source_bbox", "bbox_normalized", "canvas_to_original", "crop_to_canvas"):
|
|
self.assertEqual(stitcher[key], old_stitcher[key], key)
|
|
result = model[:, :height, :width].clone()
|
|
for seam, match in (("classic", 0.0), ("classic", 1.0), (SEAM_BLEND_IN, 0.0), (SEAM_BLEND_IN, 1.0)):
|
|
self.assertTrue(torch.equal(
|
|
apply_stitch(stitcher, result, color_match=match, seam=seam),
|
|
apply_stitch(old_stitcher, result, color_match=match, seam=seam),
|
|
), (seam, match))
|
|
old_original = torch.from_numpy(np.asarray(source.convert("RGB"), dtype=np.float32).copy() / 255.0)
|
|
self.assertTrue(torch.equal(original[0], old_original))
|
|
self.assertTrue(torch.equal(prompt_image, image))
|
|
|
|
def test_the_video_frame_node_is_unchanged(self):
|
|
for source_name, source in opaque_sources():
|
|
with self.subTest(source=source_name):
|
|
with patch("nodes.node_video_crop_rotate_pad.resolve_video_path", return_value=Path("v.mp4")), patch(
|
|
"nodes.node_video_crop_rotate_pad.decode_video_frame", return_value=(source.copy(), 0, 0)
|
|
):
|
|
outputs = AusBossVideoCropRotatePad().load_transform(
|
|
"v.mp4", "input folder", "", "frame index", 0, 0, rotation_degrees=11, pad_left=20, feather=8
|
|
)
|
|
old_output, old_mask, _ = reference_transform(source, TransformSpec(rotation_degrees=11, pad_left=20, feather=8))
|
|
self.assertEqual(len(outputs), 6)
|
|
self.assertTrue(torch.equal(outputs[0], as_tensor(old_output)))
|
|
self.assertTrue(torch.equal(outputs[1], as_tensor(old_mask)))
|
|
|
|
|
|
GREEN = (20, 190, 60)
|
|
HIDDEN = (255, 0, 0) # what a see-through pixel stores: never seen, never pasted
|
|
|
|
|
|
def cutout(width=96, height=72, alpha=None):
|
|
"""A green picture whose see-through pixels (below 90% of the most solid
|
|
one) store pure red."""
|
|
if alpha is None:
|
|
alpha = np.zeros((height, width), np.uint8)
|
|
alpha[14:58, 20:76] = 255
|
|
rgb = np.empty((height, width, 3), np.uint8)
|
|
rgb[:] = GREEN
|
|
rgb[alpha.astype(np.int32) * 100 < int(alpha.max()) * 90] = HIDDEN
|
|
return Image.fromarray(np.dstack([rgb, alpha]), "RGBA")
|
|
|
|
|
|
class SeeThroughRuleTests(unittest.TestCase):
|
|
def test_ninety_percent_of_the_most_solid_pixel_is_kept(self):
|
|
alpha = np.array([[0, 25, 128, 229, 230, 254, 255]], np.uint8)
|
|
kept = see_through_kept(Image.fromarray(np.dstack([np.zeros((1, 7, 3), np.uint8), alpha]), "RGBA"))
|
|
self.assertEqual(kept.tolist(), [[False, False, False, False, True, True, True]])
|
|
# Measured against the most solid pixel: here 200, so 180 stays.
|
|
alpha = np.array([[0, 100, 179, 180, 200]], np.uint8)
|
|
kept = see_through_kept(Image.fromarray(np.dstack([np.zeros((1, 5, 3), np.uint8), alpha]), "RGBA"))
|
|
self.assertEqual(kept.tolist(), [[False, False, False, True, True]])
|
|
|
|
def test_a_picture_without_see_through_parts_is_left_alone(self):
|
|
self.assertIsNone(see_through_kept(Image.new("RGB", (4, 4))))
|
|
self.assertIsNone(see_through_kept(Image.new("RGBA", (4, 4), (1, 2, 3, 255))))
|
|
self.assertIsNone(see_through_kept(Image.new("RGBA", (4, 4), (1, 2, 3, 230))))
|
|
self.assertIsNotNone(see_through_kept(Image.new("RGBA", (4, 4), (1, 2, 3, 229))))
|
|
|
|
def test_see_through_pixels_become_fill_and_area_to_paint(self):
|
|
alpha = np.zeros((8, 10), np.uint8)
|
|
alpha[:, 5:] = 255
|
|
alpha[0, :5] = (10, 100, 200, 229, 230) # the last one is kept
|
|
image, mask, _, _, _, _, _ = run_node(cutout(10, 8, alpha), feather=0)
|
|
see_through = torch.from_numpy(alpha < 230)
|
|
fill = torch.tensor(fill_rgb("#808080"), dtype=torch.float32) / 255
|
|
self.assertTrue(torch.equal(image[0][see_through], fill.expand(int(see_through.sum()), 3)))
|
|
self.assertTrue(bool((mask[0][see_through] == 1).all()))
|
|
kept = ~see_through
|
|
# Kept pixels are fully solid in their stored colour, 90% or 100%.
|
|
stored = torch.from_numpy(np.asarray(cutout(10, 8, alpha).convert("RGB"), dtype=np.float32) / 255)
|
|
self.assertTrue(torch.equal(image[0][kept], stored[kept]))
|
|
self.assertTrue(bool((mask[0][kept] == 0).all()))
|
|
|
|
def test_the_hidden_colour_never_reaches_the_canvas(self):
|
|
# Soft edges: the see-through ramp stores red at every alpha below 90%.
|
|
alpha = np.asarray(Image.fromarray(np.asarray(cutout().getchannel("A"))).filter(ImageFilter.GaussianBlur(6)))
|
|
source = cutout(alpha=alpha)
|
|
for name, layout in LAYOUTS.items():
|
|
with self.subTest(layout=name):
|
|
image, _, stitcher, *_ = run_node(source, **layout)
|
|
fill = fill_rgb(layout.get("fill_color", "#808080"))
|
|
# Green and the fill mix at turned edges; red above both can
|
|
# only come from the hidden colour.
|
|
ceiling = max(fill[0], GREEN[0]) + 0.5
|
|
self.assertLessEqual(float(image[..., 0].max()) * 255, ceiling)
|
|
self.assertLessEqual(float(stitcher["canvas"][..., 0].max()) * 255, ceiling)
|
|
|
|
def test_stitch_inpaint_never_pastes_see_through_parts_back(self):
|
|
alpha = np.asarray(Image.fromarray(np.asarray(cutout().getchannel("A"))).filter(ImageFilter.GaussianBlur(4)))
|
|
source = cutout(alpha=alpha)
|
|
see_through = torch.from_numpy(alpha.astype(np.int32) * 100 < 255 * 90)
|
|
self.assertTrue(bool(see_through.any()))
|
|
for blend in (0, 32):
|
|
image, mask, stitcher, *_ = run_node(source, feather=12, stitch_blend=blend)
|
|
model = torch.zeros_like(image)
|
|
model[..., 0] = 1.0
|
|
model[..., 2] = 1.0
|
|
for seam, match in (("classic", 0.0), ("classic", 1.0), (SEAM_BLEND_IN, 0.0), (SEAM_BLEND_IN, 1.0)):
|
|
with self.subTest(blend=blend, seam=seam, match=match):
|
|
out = apply_stitch(stitcher, model, color_match=match, seam=seam)[0][see_through]
|
|
if seam == SEAM_BLEND_IN and match == 0:
|
|
# New area: the model's pixels, bit for bit.
|
|
self.assertTrue(torch.equal(out, model[0][see_through]))
|
|
elif blend == 0 and match == 0:
|
|
self.assertTrue(torch.allclose(out, model[0][see_through], atol=1e-6))
|
|
elif match == 0:
|
|
# Whatever share the canvas keeps there is the fill,
|
|
# never the stored red: magenta keeps its green at 0.
|
|
self.assertLess(float(out[:, 1].max()), 0.01)
|
|
# blend in reads the see-through parts as new area.
|
|
plan = seam_plan(stitcher)
|
|
self.assertTrue(bool((plan["depth"][0][see_through] < 0).all()))
|
|
|
|
def test_what_a_see_through_pixel_stores_changes_nothing(self):
|
|
# The same cutout twice, its see-through pixels storing red in one
|
|
# and blue in the other: every output and every stitch is the same.
|
|
alpha = np.asarray(Image.fromarray(np.asarray(cutout().getchannel("A"))).filter(ImageFilter.GaussianBlur(5)))
|
|
red = cutout(alpha=alpha)
|
|
blue_rgb = np.array(red)
|
|
blue_rgb[alpha.astype(np.int32) * 100 < 255 * 90, :3] = (0, 0, 255)
|
|
blue = Image.fromarray(blue_rgb, "RGBA")
|
|
model = torch.from_numpy(np.random.default_rng(2).random((1, 72, 150, 3), dtype=np.float32))
|
|
for name, layout in LAYOUTS.items():
|
|
with self.subTest(layout=name):
|
|
a, b = run_node(red, **layout), run_node(blue, **layout)
|
|
for index in (0, 1, 3, 6):
|
|
self.assertTrue(torch.equal(a[index], b[index]), index)
|
|
self.assertTrue(torch.equal(a[2]["canvas"], b[2]["canvas"]))
|
|
result = model[:, : a[5], : a[4]].clone()
|
|
if result.shape[1:3] != a[0].shape[1:3]:
|
|
result = torch.rand_like(a[0])
|
|
for seam, match in (("classic", 1.0), (SEAM_BLEND_IN, 1.0)):
|
|
self.assertTrue(torch.equal(
|
|
apply_stitch(a[2], result, color_match=match, seam=seam),
|
|
apply_stitch(b[2], result, color_match=match, seam=seam),
|
|
))
|
|
|
|
def test_a_hole_in_the_middle_is_painted(self):
|
|
alpha = np.full((80, 100), 255, np.uint8)
|
|
yy, xx = np.mgrid[:80, :100]
|
|
hole = (yy - 40) ** 2 / 15**2 + (xx - 50) ** 2 / 22**2 <= 1
|
|
alpha[hole] = 0
|
|
image, mask, stitcher, original, _, _, prompt_image = run_node(cutout(100, 80, alpha), feather=6, stitch_blend=4)
|
|
hole_t = torch.from_numpy(hole)
|
|
self.assertTrue(bool((mask[0][hole_t] == 1).all()))
|
|
self.assertTrue(bool((prompt_image[0][hole_t] == 1).all()))
|
|
self.assertTrue(bool((original[0][hole_t] == 1).all()))
|
|
model = torch.full_like(image, 0.25)
|
|
plan = seam_plan(stitcher)
|
|
kept = {
|
|
"classic": stitcher["blend"][0] == 0,
|
|
SEAM_BLEND_IN: plan["depth"][0] >= max(plan["tone"][1], plan["detail"][1]),
|
|
}
|
|
for seam in ("classic", SEAM_BLEND_IN):
|
|
with self.subTest(seam=seam):
|
|
out = apply_stitch(stitcher, model, seam=seam)
|
|
self.assertTrue(torch.allclose(out[0][hole_t], model[0][hole_t], atol=1e-6))
|
|
# The picture around the hole comes back as it was.
|
|
self.assertTrue(bool(kept[seam].any()))
|
|
self.assertTrue(torch.equal(out[0][kept[seam]], image[0][kept[seam]]))
|
|
|
|
def test_a_turned_cutout_marks_both_kinds_of_empty(self):
|
|
source = cutout()
|
|
image, mask, stitcher, original, width, height, prompt_image = run_node(
|
|
source, rotation_degrees=17, pad_left=30, pad_right=30, feather=0
|
|
)
|
|
fill = torch.tensor(fill_rgb("#808080"), dtype=torch.float32) / 255
|
|
# The padding and the corners the turn opens keep the fill in the
|
|
# prompt view; the picture's see-through parts show white.
|
|
self.assertTrue(torch.equal(prompt_image[0, :, :30], image[0, :, :30]))
|
|
self.assertTrue(torch.equal(prompt_image[0, :, -30:], image[0, :, -30:]))
|
|
self.assertTrue(torch.allclose(prompt_image[0, 0, 32], fill, atol=0.5 / 255))
|
|
self.assertTrue(torch.allclose(prompt_image[0, -1, -33], fill, atol=0.5 / 255))
|
|
white = (prompt_image[0] == 1).all(dim=-1)
|
|
self.assertGreater(int(white.sum()), 500)
|
|
# Fully see-through pixels: fill in the canvas, area to paint.
|
|
self.assertTrue(bool(((image[0][white] - fill).abs() < 0.5 / 255).all()))
|
|
self.assertTrue(bool((mask[0][white] == 1).all()))
|
|
# The picture itself is the same in both.
|
|
solid = mask[0] == 0
|
|
self.assertTrue(torch.equal(prompt_image[0][solid], image[0][solid]))
|
|
# And the hidden red is nowhere.
|
|
self.assertLessEqual(float(image[..., 0].max()) * 255, 128.5)
|
|
self.assertLessEqual(float(prompt_image[0][~white][:, 0].max()) * 255, 255)
|
|
self.assertLess(float((prompt_image[..., 0] - prompt_image[..., 1]).max()), 0.01)
|
|
|
|
def test_a_picture_faded_throughout_is_kept_whole(self):
|
|
alpha = np.full((40, 60), 128, np.uint8)
|
|
alpha[:, :4] = 0
|
|
image, mask, _, original, _, _, prompt_image = run_node(cutout(60, 40, alpha), feather=0)
|
|
self.assertTrue(bool((mask[0][:, 4:] == 0).all()))
|
|
self.assertTrue(bool((mask[0][:, :4] == 1).all()))
|
|
green = torch.tensor(GREEN, dtype=torch.float32) / 255
|
|
self.assertTrue(torch.allclose(image[0][:, 4:], green.expand(40, 56, 3)))
|
|
self.assertTrue(torch.allclose(original[0][:, 4:], green.expand(40, 56, 3)))
|
|
self.assertTrue(bool((original[0][:, :4] == 1).all()))
|
|
|
|
def test_a_fully_see_through_picture_is_all_paint(self):
|
|
image, mask, stitcher, original, _, _, prompt_image = run_node(cutout(30, 20, np.zeros((20, 30), np.uint8)), pad_left=10)
|
|
self.assertTrue(bool((mask == 1).all()))
|
|
fill = torch.tensor(fill_rgb("#808080"), dtype=torch.float32) / 255
|
|
self.assertTrue(torch.equal(image[0], fill.expand(20, 40, 3)))
|
|
self.assertTrue(bool((prompt_image[0][:, 10:] == 1).all()))
|
|
self.assertTrue(torch.equal(prompt_image[0][:, :10], image[0][:, :10]))
|
|
self.assertTrue(bool((original == 1).all()))
|
|
model = torch.rand_like(image)
|
|
self.assertTrue(torch.equal(apply_stitch(stitcher, model, seam=SEAM_BLEND_IN), model))
|
|
|
|
def test_each_frame_follows_its_own_alpha(self):
|
|
solid = Image.new("RGBA", (40, 30), (*GREEN, 255))
|
|
alpha = np.full((30, 40), 255, np.uint8)
|
|
alpha[:, :10] = 0
|
|
image, mask, _, original, _, _, prompt_image = run_node(solid, solid, cutout(40, 30, alpha), pad_left=8)
|
|
self.assertTrue(torch.equal(prompt_image[0], image[0]))
|
|
self.assertTrue(bool((prompt_image[1][:, 8:18] == 1).all()))
|
|
self.assertTrue(bool((mask[1][:, 8:18] == 1).all()))
|
|
self.assertTrue(bool((mask[0][:, 8:] == 0).all()))
|
|
self.assertTrue(bool((original[1][:, :10] == 1).all()))
|
|
|
|
def test_a_four_channel_batch_follows_the_same_rule(self):
|
|
array = np.asarray(cutout(), dtype=np.float32) / 255
|
|
output, mask, _ = transform_tensor_batch(torch.from_numpy(array).unsqueeze(0), TransformSpec())
|
|
see_through = torch.from_numpy(np.asarray(cutout().getchannel("A")) < 230)
|
|
self.assertTrue(bool((mask[0][see_through] == 1).all()))
|
|
self.assertLessEqual(float(output[..., 0].max()) * 255, 128.5)
|
|
|
|
@unittest.skipUnless(COMFY_ROOT, "Set AUSBOSS_COMFY_ROOT for core resize integration")
|
|
def test_the_prompt_image_follows_the_resize(self):
|
|
alpha = np.full((96, 128), 255, np.uint8)
|
|
alpha[:, :40] = 0
|
|
source = cutout(128, 96, alpha)
|
|
image, mask, _, _, width, height, prompt_image = run_node(
|
|
source, pad_right=32, resize_to_megapixels=True, megapixels=0.03, resolution_steps=8
|
|
)
|
|
self.assertEqual(tuple(prompt_image.shape), tuple(image.shape))
|
|
# Well clear of the see-through band the two are the same picture.
|
|
right = int(width * 0.45)
|
|
self.assertLess(float((prompt_image[0, :, right:] - image[0, :, right:]).abs().max()), 1e-6)
|
|
self.assertGreater(float(prompt_image[0, :, : int(width * 0.2)].min()), 0.98)
|
|
|
|
|
|
class KeepPictureTests(unittest.TestCase):
|
|
"""painted_area "keep picture": the mask still marks the painted part,
|
|
and the picture under it stays. A mask painted in the MaskEditor is the
|
|
picture's alpha, with the whole picture still stored under it."""
|
|
|
|
def painted(self, width=96, height=72):
|
|
alpha = np.full((height, width), 255, np.uint8)
|
|
alpha[20:50, 30:70] = 0
|
|
return cutout(width, height, alpha), alpha
|
|
|
|
def test_the_mask_is_the_same_and_the_picture_is_whole(self):
|
|
source, alpha = self.painted()
|
|
for layout_name, layout in LAYOUTS.items():
|
|
with self.subTest(layout=layout_name):
|
|
filled = run_node(source, **layout)
|
|
kept = run_node(source, painted_area="keep picture", **layout)
|
|
solid = run_node(source.convert("RGB"), **layout)
|
|
# The mask and the size are what "fill" gives.
|
|
self.assertTrue(torch.equal(kept[1], filled[1]))
|
|
self.assertEqual(kept[4:6], filled[4:6])
|
|
# The picture is the one an opaque copy of the file gives.
|
|
self.assertTrue(torch.equal(kept[0], solid[0]))
|
|
# original and prompt_image carry the picture too, not white.
|
|
self.assertTrue(torch.equal(kept[3], solid[3]))
|
|
self.assertTrue(torch.equal(kept[6], kept[0]))
|
|
|
|
def test_fill_is_still_what_a_run_without_the_input_gives(self):
|
|
source, _ = self.painted()
|
|
default = run_node(source, pad_left=16)
|
|
named = run_node(source, pad_left=16, painted_area="fill")
|
|
for a, b in zip(default, named):
|
|
if isinstance(a, torch.Tensor):
|
|
self.assertTrue(torch.equal(a, b))
|
|
# Under the paint: the fill, not the picture.
|
|
fill = torch.tensor(fill_rgb("#808080"), dtype=torch.float32) / 255.0
|
|
self.assertTrue(torch.allclose(default[0][0, 30, 16 + 40], fill, atol=1e-6))
|
|
|
|
def test_stitch_inpaint_still_takes_the_painted_part_from_the_result(self):
|
|
source, alpha = self.painted()
|
|
image, mask, stitcher, *_ = run_node(source, painted_area="keep picture", stitch_blend=0)
|
|
result = torch.zeros_like(image)
|
|
stitched = apply_stitch(stitcher, result, color_match=0.0)
|
|
# Painted part: the result. Well outside it: the picture.
|
|
self.assertEqual(float(stitched[0, 35, 50].abs().max()), 0.0)
|
|
self.assertTrue(torch.equal(stitched[0, 5, 5], image[0, 5, 5]))
|
|
|
|
def test_a_picture_without_see_through_parts_is_unchanged(self):
|
|
for source_name, source in opaque_sources():
|
|
with self.subTest(source=source_name):
|
|
plain = run_node(source, pad_top=8, rotation_degrees=7.0)
|
|
kept = run_node(source, pad_top=8, rotation_degrees=7.0, painted_area="keep picture")
|
|
for a, b in zip(plain, kept):
|
|
if isinstance(a, torch.Tensor):
|
|
self.assertTrue(torch.equal(a, b))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|