mirror of
https://github.com/Bing-su/adetailer.git
synced 2026-01-26 11:19:53 +00:00
233 lines
6.8 KiB
Python
233 lines
6.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image, ImageDraw
|
|
|
|
from adetailer.mask import (
|
|
bbox_area,
|
|
dilate_erode,
|
|
has_intersection,
|
|
is_all_black,
|
|
mask_invert,
|
|
mask_merge,
|
|
offset,
|
|
)
|
|
|
|
|
|
def test_dilate_positive_value():
|
|
img = Image.new("L", (10, 10), color="black")
|
|
draw = ImageDraw.Draw(img)
|
|
draw.rectangle((3, 3, 5, 5), fill="white")
|
|
value = 3
|
|
|
|
result = dilate_erode(img, value)
|
|
|
|
assert isinstance(result, Image.Image)
|
|
assert result.size == (10, 10)
|
|
|
|
expect = np.array(
|
|
[
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
|
|
[0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
|
|
[0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
|
|
[0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
|
|
[0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
],
|
|
dtype=np.uint8,
|
|
)
|
|
assert np.array_equal(np.array(result), expect)
|
|
|
|
|
|
def test_offset():
|
|
img = Image.new("L", (10, 10), color="black")
|
|
draw = ImageDraw.Draw(img)
|
|
draw.rectangle((4, 4, 5, 5), fill="white")
|
|
|
|
result = offset(img, x=1, y=2)
|
|
|
|
assert isinstance(result, Image.Image)
|
|
assert result.size == (10, 10)
|
|
|
|
expect = np.array(
|
|
[
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 255, 255, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 255, 255, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
],
|
|
dtype=np.uint8,
|
|
)
|
|
assert np.array_equal(np.array(result), expect)
|
|
|
|
|
|
class TestIsAllBlack:
|
|
def test_is_all_black_1(self):
|
|
img = Image.new("L", (10, 10), color="black")
|
|
assert is_all_black(img)
|
|
|
|
draw = ImageDraw.Draw(img)
|
|
draw.rectangle((4, 4, 5, 5), fill="white")
|
|
assert not is_all_black(img)
|
|
|
|
def test_is_all_black_2(self):
|
|
img = np.zeros((10, 10), dtype=np.uint8)
|
|
assert is_all_black(img)
|
|
|
|
img[4:6, 4:6] = 255
|
|
assert not is_all_black(img)
|
|
|
|
def test_is_all_black_rgb_image_pil(self):
|
|
img = Image.new("RGB", (10, 10), color="red")
|
|
assert not is_all_black(img)
|
|
|
|
img = Image.new("RGBA", (10, 10), color="red")
|
|
assert not is_all_black(img)
|
|
|
|
def test_is_all_black_rgb_image_numpy(self):
|
|
img = np.full((10, 10, 4), 127, dtype=np.uint8)
|
|
with pytest.raises(cv2.error):
|
|
is_all_black(img)
|
|
|
|
img = np.full((4, 10, 10), 0.5, dtype=np.float32)
|
|
with pytest.raises(cv2.error):
|
|
is_all_black(img)
|
|
|
|
|
|
class TestHasIntersection:
|
|
def test_has_intersection_1(self):
|
|
arr1 = np.array(
|
|
[
|
|
[0, 0, 0, 0],
|
|
[0, 0, 0, 0],
|
|
[0, 0, 0, 0],
|
|
[0, 0, 0, 0],
|
|
]
|
|
)
|
|
arr2 = arr1.copy()
|
|
assert not has_intersection(arr1, arr2)
|
|
|
|
def test_has_intersection_2(self):
|
|
arr1 = np.array(
|
|
[
|
|
[0, 0, 0, 0],
|
|
[0, 255, 255, 0],
|
|
[0, 255, 255, 0],
|
|
[0, 0, 0, 0],
|
|
]
|
|
)
|
|
arr2 = np.array(
|
|
[
|
|
[0, 0, 0, 0],
|
|
[0, 0, 0, 0],
|
|
[0, 0, 255, 255],
|
|
[0, 0, 255, 255],
|
|
]
|
|
)
|
|
assert has_intersection(arr1, arr2)
|
|
|
|
arr3 = np.array(
|
|
[
|
|
[255, 0, 0, 0],
|
|
[0, 0, 0, 0],
|
|
[0, 0, 0, 255],
|
|
[0, 0, 255, 255],
|
|
]
|
|
)
|
|
assert not has_intersection(arr1, arr3)
|
|
|
|
def test_has_intersection_3(self):
|
|
img1 = Image.new("L", (10, 10), color="black")
|
|
draw1 = ImageDraw.Draw(img1)
|
|
draw1.rectangle((3, 3, 5, 5), fill="white")
|
|
img2 = Image.new("L", (10, 10), color="black")
|
|
draw2 = ImageDraw.Draw(img2)
|
|
draw2.rectangle((6, 6, 8, 8), fill="white")
|
|
assert not has_intersection(img1, img2)
|
|
|
|
img3 = Image.new("L", (10, 10), color="black")
|
|
draw3 = ImageDraw.Draw(img3)
|
|
draw3.rectangle((2, 2, 8, 8), fill="white")
|
|
assert has_intersection(img1, img3)
|
|
|
|
def test_has_intersection_4(self):
|
|
img1 = Image.new("RGB", (10, 10), color="black")
|
|
draw1 = ImageDraw.Draw(img1)
|
|
draw1.rectangle((3, 3, 5, 5), fill="white")
|
|
img2 = Image.new("RGBA", (10, 10), color="black")
|
|
draw2 = ImageDraw.Draw(img2)
|
|
draw2.rectangle((2, 2, 8, 8), fill="white")
|
|
assert has_intersection(img1, img2)
|
|
|
|
def test_has_intersection_5(self):
|
|
img1 = Image.new("RGB", (10, 10), color="black")
|
|
draw1 = ImageDraw.Draw(img1)
|
|
draw1.rectangle((4, 4, 5, 5), fill="white")
|
|
img2 = np.full((10, 10, 4), 255, dtype=np.uint8)
|
|
assert has_intersection(img1, img2)
|
|
|
|
|
|
def test_bbox_area():
|
|
bbox = [0.0, 0.0, 10.0, 10.0]
|
|
assert bbox_area(bbox) == 100
|
|
|
|
|
|
class TestMaskMerge:
|
|
def test_mask_merge(self):
|
|
img1 = Image.new("L", (10, 10), color="black")
|
|
draw1 = ImageDraw.Draw(img1)
|
|
draw1.rectangle((3, 3, 5, 5), fill="white")
|
|
|
|
img2 = Image.new("L", (10, 10), color="black")
|
|
draw2 = ImageDraw.Draw(img2)
|
|
draw2.rectangle((6, 6, 8, 8), fill="white")
|
|
|
|
merged = mask_merge([img1, img2])
|
|
assert len(merged) == 1
|
|
|
|
expect = Image.new("L", (10, 10), color="black")
|
|
draw3 = ImageDraw.Draw(expect)
|
|
draw3.rectangle((3, 3, 5, 5), fill="white")
|
|
draw3.rectangle((6, 6, 8, 8), fill="white")
|
|
|
|
assert np.array_equal(np.array(merged[0]), np.array(expect))
|
|
|
|
def test_merge_mask_different_size(self):
|
|
img1 = Image.new("L", (10, 10), color="black")
|
|
draw1 = ImageDraw.Draw(img1)
|
|
draw1.rectangle((3, 3, 5, 5), fill="white")
|
|
|
|
img2 = Image.new("L", (20, 20), color="black")
|
|
draw2 = ImageDraw.Draw(img2)
|
|
draw2.rectangle((6, 6, 8, 8), fill="white")
|
|
|
|
with pytest.raises(
|
|
cv2.error, match="-209:Sizes of input arguments do not match"
|
|
):
|
|
mask_merge([img1, img2])
|
|
|
|
|
|
def test_mask_invert():
|
|
img = Image.new("L", (10, 10), color="black")
|
|
draw = ImageDraw.Draw(img)
|
|
draw.rectangle((3, 3, 5, 5), fill="white")
|
|
|
|
inverted = mask_invert([img])
|
|
assert len(inverted) == 1
|
|
|
|
expect = Image.new("L", (10, 10), color="white")
|
|
draw = ImageDraw.Draw(expect)
|
|
draw.rectangle((3, 3, 5, 5), fill="black")
|
|
|
|
assert np.array_equal(np.array(inverted[0]), np.array(expect))
|