File size: 2,344 Bytes
3dabe4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import numpy as np
from PIL import Image, ImageDraw
from adetailer.common import create_bbox_from_mask, create_mask_from_bbox
def test_create_mask_from_bbox():
img = Image.new("L", (10, 10), color="black")
bbox = [[1.0, 1.0, 2.0, 2.0], [7.0, 7.0, 8.0, 8.0]]
masks = create_mask_from_bbox(bbox, img.size)
expect1 = np.array(
[
[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, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
],
dtype=np.uint8,
)
expect2 = 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, 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, 255, 255, 0],
[0, 0, 0, 0, 0, 0, 0, 255, 255, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
],
dtype=np.uint8,
)
assert len(masks) == len(bbox)
arr1 = np.array(masks[0])
arr2 = np.array(masks[1])
assert arr1.shape == expect1.shape
assert arr2.shape == expect2.shape
assert arr1.shape == (10, 10)
assert arr1.dtype == expect1.dtype
assert arr2.dtype == expect2.dtype
assert np.array_equal(arr1, expect1)
assert np.array_equal(arr2, expect2)
# The function correctly receives a list of masks and the shape of the image.
def test_create_bbox_from_mask():
mask = Image.new("L", (10, 10), color="black")
draw = ImageDraw.Draw(mask)
draw.rectangle((2, 2, 5, 5), fill="white")
result = create_bbox_from_mask([mask], (10, 10))
assert isinstance(result, list)
assert len(result) == 1
assert all(isinstance(bbox, list) for bbox in result)
assert all(len(bbox) == 4 for bbox in result)
assert result[0] == [2, 2, 6, 6]
result = create_bbox_from_mask([mask], (256, 256))
assert result[0] == [38, 38, 166, 166]
|