2023-09-13 22:27:44 +00:00
|
|
|
import unittest
|
|
|
|
|
|
|
|
from PIL import Image
|
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
from onnx_web.chain.result import StageResult
|
2023-09-15 00:35:48 +00:00
|
|
|
from onnx_web.chain.tile import (
|
|
|
|
complete_tile,
|
2023-11-20 00:39:39 +00:00
|
|
|
generate_tile_grid,
|
2023-09-15 00:35:48 +00:00
|
|
|
generate_tile_spiral,
|
2023-11-20 00:39:39 +00:00
|
|
|
make_tile_grads,
|
2023-09-15 00:35:48 +00:00
|
|
|
needs_tile,
|
2023-11-20 00:39:39 +00:00
|
|
|
process_tile_stack,
|
2023-09-15 00:35:48 +00:00
|
|
|
)
|
2023-09-14 03:03:39 +00:00
|
|
|
from onnx_web.params import Size
|
2023-09-13 22:27:44 +00:00
|
|
|
|
|
|
|
|
|
|
|
class TestCompleteTile(unittest.TestCase):
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_complete_tile(self):
|
|
|
|
partial = Image.new("RGB", (64, 64))
|
|
|
|
output = complete_tile(partial, 64)
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(output.size, (64, 64))
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_partial_tile(self):
|
|
|
|
partial = Image.new("RGB", (64, 32))
|
|
|
|
output = complete_tile(partial, 64)
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(output.size, (64, 64))
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_nothing(self):
|
|
|
|
output = complete_tile(None, 64)
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertIsNone(output)
|
2023-09-13 22:27:44 +00:00
|
|
|
|
|
|
|
|
|
|
|
class TestNeedsTile(unittest.TestCase):
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_undersized_source(self):
|
|
|
|
small = Image.new("RGB", (32, 32))
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertFalse(needs_tile(64, 64, source=small))
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_oversized_source(self):
|
|
|
|
large = Image.new("RGB", (64, 64))
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertTrue(needs_tile(32, 32, source=large))
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_undersized_size(self):
|
|
|
|
small = Size(32, 32)
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertFalse(needs_tile(64, 64, size=small))
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_oversized_size(self):
|
|
|
|
large = Size(64, 64)
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertTrue(needs_tile(32, 32, size=large))
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_with_nothing(self):
|
|
|
|
self.assertFalse(needs_tile(32, 32))
|
2023-09-13 22:27:44 +00:00
|
|
|
|
|
|
|
|
|
|
|
class TestTileGrads(unittest.TestCase):
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_center_tile(self):
|
|
|
|
grad_x, grad_y = make_tile_grads(32, 32, 8, 64, 64)
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(grad_x, [0, 1, 1, 0])
|
|
|
|
self.assertEqual(grad_y, [0, 1, 1, 0])
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_vertical_edge_tile(self):
|
|
|
|
grad_x, grad_y = make_tile_grads(32, 0, 8, 64, 8)
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(grad_x, [0, 1, 1, 0])
|
|
|
|
self.assertEqual(grad_y, [1, 1, 1, 1])
|
2023-09-14 03:03:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_horizontal_edge_tile(self):
|
|
|
|
grad_x, grad_y = make_tile_grads(0, 32, 8, 8, 64)
|
2023-09-13 22:27:44 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(grad_x, [1, 1, 1, 1])
|
|
|
|
self.assertEqual(grad_y, [0, 1, 1, 0])
|
2023-09-14 03:40:22 +00:00
|
|
|
|
|
|
|
|
2023-11-20 00:39:39 +00:00
|
|
|
class TestGenerateTileGrid(unittest.TestCase):
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_grid_complete(self):
|
|
|
|
tiles = generate_tile_grid(16, 16, 8, 0.0)
|
2023-11-20 00:39:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(len(tiles), 4)
|
|
|
|
self.assertEqual(tiles, [(0, 0), (8, 0), (0, 8), (8, 8)])
|
2023-11-20 00:39:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_grid_no_overlap(self):
|
|
|
|
tiles = generate_tile_grid(64, 64, 8, 0.0)
|
2023-11-20 00:39:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(len(tiles), 64)
|
|
|
|
self.assertEqual(tiles[0:4], [(0, 0), (8, 0), (16, 0), (24, 0)])
|
|
|
|
self.assertEqual(tiles[-5:-1], [(24, 56), (32, 56), (40, 56), (48, 56)])
|
2023-11-20 00:39:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_grid_50_overlap(self):
|
|
|
|
tiles = generate_tile_grid(64, 64, 8, 0.5)
|
2023-11-20 00:39:39 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(len(tiles), 256)
|
|
|
|
self.assertEqual(tiles[0:4], [(0, 0), (4, 0), (8, 0), (12, 0)])
|
|
|
|
self.assertEqual(tiles[-5:-1], [(44, 60), (48, 60), (52, 60), (56, 60)])
|
2023-11-20 00:39:39 +00:00
|
|
|
|
|
|
|
|
2023-09-14 03:40:22 +00:00
|
|
|
class TestGenerateTileSpiral(unittest.TestCase):
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_spiral_complete(self):
|
|
|
|
tiles = generate_tile_spiral(16, 16, 8, 0.0)
|
2023-09-14 03:40:22 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(len(tiles), 4)
|
|
|
|
self.assertEqual(tiles, [(0, 0), (8, 0), (8, 8), (0, 8)])
|
2023-09-14 03:40:22 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_spiral_no_overlap(self):
|
|
|
|
tiles = generate_tile_spiral(64, 64, 8, 0.0)
|
2023-09-14 03:40:22 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(len(tiles), 64)
|
|
|
|
self.assertEqual(tiles[0:4], [(0, 0), (8, 0), (16, 0), (24, 0)])
|
|
|
|
self.assertEqual(tiles[-5:-1], [(16, 24), (24, 24), (32, 24), (32, 32)])
|
2023-09-14 03:40:22 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_spiral_50_overlap(self):
|
|
|
|
tiles = generate_tile_spiral(64, 64, 8, 0.5)
|
2023-09-14 03:40:22 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(len(tiles), 225)
|
|
|
|
self.assertEqual(tiles[0:4], [(0, 0), (4, 0), (8, 0), (12, 0)])
|
|
|
|
self.assertEqual(tiles[-5:-1], [(32, 32), (28, 32), (24, 32), (24, 28)])
|
2023-09-28 23:45:04 +00:00
|
|
|
|
|
|
|
|
2023-11-20 00:39:39 +00:00
|
|
|
class TestProcessTileStack(unittest.TestCase):
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_grid_full(self):
|
|
|
|
source = Image.new("RGB", (64, 64))
|
|
|
|
blend = process_tile_stack(
|
|
|
|
StageResult(images=[source]), 32, 1, [], generate_tile_grid
|
|
|
|
)
|
2023-09-28 23:45:04 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(blend[0].size, (64, 64))
|
2023-09-28 23:45:04 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
def test_grid_partial(self):
|
|
|
|
source = Image.new("RGB", (72, 72))
|
|
|
|
blend = process_tile_stack(
|
|
|
|
StageResult(images=[source]), 32, 1, [], generate_tile_grid
|
|
|
|
)
|
2023-09-28 23:45:04 +00:00
|
|
|
|
2023-11-20 05:18:57 +00:00
|
|
|
self.assertEqual(blend[0].size, (72, 72))
|