90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
import unittest
|
|
import numpy as np
|
|
import torch
|
|
from nodes.nodes import MaskContourProcessor
|
|
|
|
class TestMaskContourProcessor(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test fixtures before each test method."""
|
|
self.processor = MaskContourProcessor()
|
|
|
|
# Create some common test masks
|
|
self.empty_mask = np.zeros((64, 64), dtype=np.float32)
|
|
self.full_mask = np.ones((64, 64), dtype=np.float32)
|
|
self.circle_mask = np.zeros((64, 64), dtype=np.float32)
|
|
center = (32, 32)
|
|
radius = 20
|
|
y, x = np.ogrid[:64, :64]
|
|
self.circle_mask[(x - center[0])**2 + (y - center[1])**2 <= radius**2] = 1
|
|
|
|
def test_calculate_mask_centroid(self):
|
|
"""Test centroid calculation for various masks."""
|
|
# Test empty mask
|
|
center_x, center_y = self.processor.calculate_mask_centroid(self.empty_mask)
|
|
self.assertEqual((center_x, center_y), (32, 32)) # Should return center of image
|
|
|
|
# Test full mask
|
|
center_x, center_y = self.processor.calculate_mask_centroid(self.full_mask)
|
|
self.assertEqual((center_x, center_y), (31.5, 31.5)) # Should be center of image
|
|
|
|
# Test circle mask
|
|
center_x, center_y = self.processor.calculate_mask_centroid(self.circle_mask)
|
|
self.assertAlmostEqual(center_x, 32, delta=1)
|
|
self.assertAlmostEqual(center_y, 32, delta=1)
|
|
|
|
def test_detect_edge_points(self):
|
|
"""Test edge point detection."""
|
|
center = (32, 32)
|
|
edge_points = self.processor.detect_edge_points(self.circle_mask, center)
|
|
|
|
# Check that we got some edge points
|
|
self.assertTrue(len(edge_points) > 0)
|
|
|
|
# Check that all points are on the edge (should have at least one
|
|
# neighboring pixel with different value)
|
|
for x, y in edge_points:
|
|
has_different_neighbor = False
|
|
for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]:
|
|
nx, ny = x + dx, y + dy
|
|
if (0 <= nx < 64 and 0 <= ny < 64 and
|
|
self.circle_mask[ny, nx] != self.circle_mask[y, x]):
|
|
has_different_neighbor = True
|
|
break
|
|
self.assertTrue(has_different_neighbor)
|
|
|
|
def test_redistribute_points(self):
|
|
"""Test point redistribution."""
|
|
# Create a simple square of points
|
|
square_points = [(0,0), (0,10), (10,10), (10,0)]
|
|
|
|
# Test redistribution to different numbers of points
|
|
for target_count in [4, 8, 16]:
|
|
new_points = self.processor.redistribute_points(square_points, target_count)
|
|
self.assertEqual(len(new_points), target_count)
|
|
|
|
def test_process_mask(self):
|
|
"""Test the main processing function."""
|
|
# Convert numpy mask to tensor
|
|
mask_tensor = torch.from_numpy(self.circle_mask)
|
|
|
|
# Process the mask
|
|
result = self.processor.process_mask(
|
|
mask_tensor,
|
|
line_length=1.0,
|
|
line_count=20,
|
|
line_width=0.01
|
|
)
|
|
|
|
# Check return type and shape
|
|
self.assertIsInstance(result, tuple)
|
|
self.assertEqual(len(result), 1)
|
|
self.assertIsInstance(result[0], torch.Tensor)
|
|
self.assertEqual(result[0].shape, mask_tensor.shape)
|
|
|
|
# Check that the output contains at least the original mask
|
|
# (output should be >= input everywhere)
|
|
self.assertTrue(torch.all(result[0] >= mask_tensor))
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|