235 lines
10 KiB
Python
235 lines
10 KiB
Python
import unittest
|
|
import torch
|
|
import numpy as np
|
|
from inpaint_cropandstitch import CPUProcessorLogic
|
|
|
|
|
|
class TestCPUProcessorLogic(unittest.TestCase):
|
|
def setUp(self):
|
|
self.processor = CPUProcessorLogic()
|
|
|
|
def test_rescale_i_algorithms(self):
|
|
img = torch.rand(1, 32, 32, 3, dtype=torch.float32)
|
|
algorithms = ["bicubic", "bilinear", "nearest", "nearest-exact", "lanczos", "box", "area", "hamming"]
|
|
for algo in algorithms:
|
|
with self.subTest(algorithm=algo):
|
|
res = self.processor.rescale_i(img, 64, 48, algo)
|
|
self.assertEqual(res.shape, (1, 48, 64, 3))
|
|
self.assertEqual(res.dtype, torch.float32)
|
|
|
|
def test_rescale_i_channels_and_batches(self):
|
|
for c in [1, 3, 4]:
|
|
with self.subTest(channels=c):
|
|
img = torch.rand(2, 20, 20, c)
|
|
res = self.processor.rescale_i(img, 40, 30, "bilinear")
|
|
self.assertEqual(res.shape, (2, 30, 40, c))
|
|
|
|
def test_rescale_i_dtypes(self):
|
|
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
|
with self.subTest(dtype=dtype):
|
|
img = torch.rand(1, 16, 16, 3, dtype=dtype)
|
|
res = self.processor.rescale_i(img, 32, 32, "bicubic")
|
|
self.assertEqual(res.shape, (1, 32, 32, 3))
|
|
self.assertIn(res.dtype, [torch.float32, dtype])
|
|
|
|
def test_rescale_i_boundary_sizes(self):
|
|
img = torch.rand(1, 10, 10, 3)
|
|
res_1x1 = self.processor.rescale_i(img, 1, 1, "bilinear")
|
|
self.assertEqual(res_1x1.shape, (1, 1, 1, 3))
|
|
|
|
res_zero = self.processor.rescale_i(img, 0, -5, "bilinear")
|
|
self.assertEqual(res_zero.shape, (1, 1, 1, 3))
|
|
|
|
def test_rescale_m_algorithms_and_shapes(self):
|
|
mask_3d = torch.rand(1, 24, 24, dtype=torch.float32)
|
|
mask_2d = torch.rand(24, 24, dtype=torch.float32)
|
|
algorithms = ["bicubic", "bilinear", "nearest", "nearest-exact", "lanczos", "box", "area", "hamming"]
|
|
for algo in algorithms:
|
|
with self.subTest(algorithm=algo):
|
|
res_3d = self.processor.rescale_m(mask_3d, 48, 36, algo)
|
|
self.assertEqual(res_3d.shape, (1, 36, 48))
|
|
res_2d = self.processor.rescale_m(mask_2d, 48, 36, algo)
|
|
self.assertEqual(res_2d.shape, (1, 36, 48))
|
|
|
|
def test_rescale_m_dtypes(self):
|
|
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
|
with self.subTest(dtype=dtype):
|
|
mask = torch.rand(1, 16, 16, dtype=dtype)
|
|
res = self.processor.rescale_m(mask, 32, 32, "bilinear")
|
|
self.assertEqual(res.shape, (1, 32, 32))
|
|
|
|
def test_fillholes_functional_hollow_ring(self):
|
|
# Frame enclosing empty hole in center
|
|
mask = torch.zeros(1, 40, 40, dtype=torch.float32)
|
|
mask[:, 5:35, 5:10] = 1.0
|
|
mask[:, 5:35, 30:35] = 1.0
|
|
mask[:, 5:10, 5:35] = 1.0
|
|
mask[:, 30:35, 5:35] = 1.0
|
|
|
|
# Center is originally zero
|
|
self.assertEqual(mask[0, 20, 20].item(), 0.0)
|
|
filled = self.processor.fillholes_iterative_hipass_fill_m(mask)
|
|
# Inside center must now be filled to 1.0
|
|
self.assertAlmostEqual(filled[0, 20, 20].item(), 1.0)
|
|
# Outside corner remains 0.0
|
|
self.assertEqual(filled[0, 0, 0].item(), 0.0)
|
|
|
|
def test_fillholes_multilevel_gradient(self):
|
|
# Soft threshold ring (0.8) with 0.0 hole
|
|
mask = torch.zeros(1, 30, 30, dtype=torch.float32)
|
|
mask[:, 5:25, 5:25] = 0.8
|
|
mask[:, 10:20, 10:20] = 0.0
|
|
filled = self.processor.fillholes_iterative_hipass_fill_m(mask)
|
|
self.assertAlmostEqual(filled[0, 15, 15].item(), 0.8, places=2)
|
|
self.assertEqual(filled[0, 0, 0].item(), 0.0)
|
|
|
|
def test_fillholes_2d_mask(self):
|
|
mask_2d = torch.zeros(40, 40, dtype=torch.float32)
|
|
mask_2d[5:35, 5:10] = 1.0
|
|
mask_2d[5:35, 30:35] = 1.0
|
|
mask_2d[5:10, 5:35] = 1.0
|
|
mask_2d[30:35, 5:35] = 1.0
|
|
filled_2d = self.processor.fillholes_iterative_hipass_fill_m(mask_2d)
|
|
self.assertEqual(filled_2d.ndim, 2)
|
|
self.assertAlmostEqual(filled_2d[20, 20].item(), 1.0)
|
|
self.assertEqual(filled_2d[0, 0].item(), 0.0)
|
|
|
|
def test_fillholes_solid_and_empty(self):
|
|
# Solid mask should remain solid
|
|
solid = torch.ones(1, 20, 20, dtype=torch.float32)
|
|
self.assertTrue(torch.equal(self.processor.fillholes_iterative_hipass_fill_m(solid), solid))
|
|
|
|
# Empty mask should remain empty
|
|
empty = torch.zeros(1, 20, 20, dtype=torch.float32)
|
|
self.assertTrue(torch.equal(self.processor.fillholes_iterative_hipass_fill_m(empty), empty))
|
|
|
|
def test_hipassfilter_m(self):
|
|
mask = torch.tensor([[[0.1, 0.4], [0.6, 0.9]]], dtype=torch.float32)
|
|
filtered = self.processor.hipassfilter_m(mask, 0.5)
|
|
self.assertEqual(filtered[0, 0, 0].item(), 0.0)
|
|
self.assertEqual(filtered[0, 0, 1].item(), 0.0)
|
|
self.assertAlmostEqual(filtered[0, 1, 0].item(), 0.6)
|
|
self.assertAlmostEqual(filtered[0, 1, 1].item(), 0.9)
|
|
|
|
filtered_neg = self.processor.hipassfilter_m(mask, -1.0)
|
|
self.assertTrue(torch.equal(filtered_neg, mask))
|
|
|
|
filtered_high = self.processor.hipassfilter_m(mask, 1.5)
|
|
self.assertEqual(torch.count_nonzero(filtered_high).item(), 0)
|
|
|
|
def test_expand_m_exact_radii(self):
|
|
mask = torch.zeros(1, 31, 31, dtype=torch.float32)
|
|
mask[0, 15, 15] = 1.0
|
|
exp = self.processor.expand_m(mask, 4)
|
|
self.assertEqual(exp[0, 15, 15].item(), 1.0)
|
|
self.assertEqual(exp[0, 15, 16].item(), 1.0)
|
|
self.assertEqual(exp[0, 15, 14].item(), 1.0)
|
|
self.assertEqual(exp[0, 0, 0].item(), 0.0)
|
|
|
|
def test_expand_m_boundary_cases(self):
|
|
mask = torch.zeros(1, 20, 20, dtype=torch.float32)
|
|
mask[0, 10, 10] = 1.0
|
|
self.assertTrue(torch.equal(self.processor.expand_m(mask, 0), mask))
|
|
self.assertTrue(torch.equal(self.processor.expand_m(mask, -5), mask))
|
|
|
|
# 2D mask
|
|
mask_2d = torch.zeros(20, 20, dtype=torch.float32)
|
|
mask_2d[10, 10] = 1.0
|
|
exp_2d = self.processor.expand_m(mask_2d, 4)
|
|
self.assertEqual(exp_2d.ndim, 2)
|
|
|
|
# Huge expansion
|
|
exp_large = self.processor.expand_m(mask, 100)
|
|
self.assertEqual(exp_large.shape, mask.shape)
|
|
|
|
def test_invert_m(self):
|
|
mask = torch.tensor([[[0.0, 0.25], [0.75, 1.0]]], dtype=torch.float32)
|
|
inv = self.processor.invert_m(mask)
|
|
self.assertAlmostEqual(inv[0, 0, 0].item(), 1.0)
|
|
self.assertAlmostEqual(inv[0, 0, 1].item(), 0.75)
|
|
self.assertAlmostEqual(inv[0, 1, 0].item(), 0.25)
|
|
self.assertAlmostEqual(inv[0, 1, 1].item(), 0.0)
|
|
|
|
# Bool mask
|
|
mask_bool = torch.tensor([[[False, True], [True, False]]], dtype=torch.bool)
|
|
inv_bool = self.processor.invert_m(mask_bool)
|
|
self.assertAlmostEqual(inv_bool[0, 0, 0].item(), 1.0)
|
|
self.assertAlmostEqual(inv_bool[0, 0, 1].item(), 0.0)
|
|
|
|
def test_blur_m_bool_and_int_dtypes(self):
|
|
mask_bool = torch.zeros(1, 21, 21, dtype=torch.bool)
|
|
mask_bool[0, 10, 10] = True
|
|
blurred_bool = self.processor.blur_m(mask_bool, 3)
|
|
self.assertGreater(blurred_bool[0, 10, 10].item(), 0.0)
|
|
|
|
mask_u8 = torch.zeros(1, 21, 21, dtype=torch.uint8)
|
|
mask_u8[0, 10, 10] = 255
|
|
blurred_u8 = self.processor.blur_m(mask_u8, 3)
|
|
self.assertGreater(blurred_u8[0, 10, 10].item(), 0.0)
|
|
|
|
def test_blur_m_gaussian_decay(self):
|
|
mask = torch.zeros(1, 31, 31, dtype=torch.float32)
|
|
mask[0, 15, 15] = 1.0
|
|
blurred = self.processor.blur_m(mask, 4)
|
|
|
|
peak = blurred[0, 15, 15].item()
|
|
d1 = blurred[0, 15, 16].item()
|
|
d2 = blurred[0, 15, 17].item()
|
|
self.assertGreater(peak, d1)
|
|
self.assertGreater(d1, d2)
|
|
self.assertAlmostEqual(blurred[0, 0, 0].item(), 0.0, places=3)
|
|
|
|
def test_pad_to_multiple(self):
|
|
for val, mult, expected in [
|
|
(0, 8, 0), (7, 8, 8), (8, 8, 8), (9, 8, 16),
|
|
(60, 8, 64), (64, 8, 64), (65, 8, 72),
|
|
(64, 16, 64), (65, 16, 80), (100, 32, 128),
|
|
(50, 0, 50), (50, -8, 50)
|
|
]:
|
|
with self.subTest(val=val, mult=mult):
|
|
self.assertEqual(self.processor.pad_to_multiple(val, mult), expected)
|
|
|
|
def test_debug_context_location_in_image_inversion(self):
|
|
img = torch.full((1, 30, 30, 3), 0.25, dtype=torch.float32)
|
|
deb = self.processor.debug_context_location_in_image(img, 10, 10, 10, 10)
|
|
# Inside box: 1.0 - 0.25 = 0.75
|
|
self.assertAlmostEqual(deb[0, 15, 15, 0].item(), 0.75)
|
|
# Outside box: remains 0.25
|
|
self.assertAlmostEqual(deb[0, 0, 0, 0].item(), 0.25)
|
|
|
|
# Clamping out-of-bounds coordinates
|
|
deb_out = self.processor.debug_context_location_in_image(img, -10, -10, 20, 20)
|
|
self.assertEqual(deb_out.shape, img.shape)
|
|
deb_far = self.processor.debug_context_location_in_image(img, 100, 100, 20, 20)
|
|
self.assertEqual(deb_far.shape, img.shape)
|
|
|
|
def test_extend_imm_edge_preservation(self):
|
|
img = torch.zeros(1, 10, 10, 3, dtype=torch.float32)
|
|
img[:, 0, :, 0] = 0.42
|
|
img[:, -1, :, 1] = 0.77
|
|
mask = torch.zeros(1, 10, 10, dtype=torch.float32)
|
|
e_img, e_mask, _ = self.processor.extend_imm(img, mask, None, 2.0, 2.0, 1.0, 1.0)
|
|
self.assertAlmostEqual(e_img[0, 0, 5, 0].item(), 0.42)
|
|
self.assertAlmostEqual(e_img[0, -1, 5, 1].item(), 0.77)
|
|
self.assertAlmostEqual(e_mask[0, 0, 5].item(), 1.0)
|
|
self.assertAlmostEqual(e_mask[0, -1, 5].item(), 1.0)
|
|
|
|
def test_preresize_imm_modes(self):
|
|
img = torch.rand(1, 100, 100, 3, dtype=torch.float32)
|
|
mask = torch.zeros(1, 100, 100, dtype=torch.float32)
|
|
opt_mask = torch.zeros(1, 100, 100, dtype=torch.float32)
|
|
|
|
# Ensure min resolution
|
|
p_img, p_mask, p_opt = self.processor.preresize_imm(img, mask, opt_mask, "bilinear", "bicubic", "ensure minimum resolution", 200, 150, 400, 400)
|
|
self.assertGreaterEqual(p_img.shape[2], 200)
|
|
self.assertGreaterEqual(p_img.shape[1], 150)
|
|
|
|
# Ensure max resolution
|
|
p_img2, p_mask2, p_opt2 = self.processor.preresize_imm(img, mask, opt_mask, "bilinear", "bicubic", "ensure maximum resolution", 10, 10, 50, 80)
|
|
self.assertLessEqual(p_img2.shape[2], 50)
|
|
self.assertLessEqual(p_img2.shape[1], 80)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|