1import numpy as np 2from skimage.graph._graph import pixel_graph, central_pixel 3 4mask = np.array([[1, 0, 0], [0, 1, 1], [0, 1, 0]], dtype=bool) 5image = np.random.default_rng().random(mask.shape) 6 7 8def test_small_graph(): 9 g, n = pixel_graph(mask, connectivity=2) 10 assert g.shape == (4, 4) 11 assert len(g.data) == 8 12 np.testing.assert_allclose(np.unique(g.data), [1, np.sqrt(2)]) 13 np.testing.assert_array_equal(n, [0, 4, 5, 7]) 14 15 16def test_central_pixel(): 17 g, n = pixel_graph(mask, connectivity=2) 18 px, ds = central_pixel(g, n, shape=mask.shape) 19 np.testing.assert_array_equal(px, (1, 1)) 20 s2 = np.sqrt(2) 21 np.testing.assert_allclose(ds, [s2*3 + 2, s2 + 2, s2*2 + 2, s2*2 + 2]) 22 23 # test raveled coordinate 24 px, _ = central_pixel(g, n) 25 assert px == 4 26 27 # test no nodes given 28 px, _ = central_pixel(g) 29 assert px == 1 30 31 32def test_edge_function(): 33 def edge_func(values_src, values_dst, distances): 34 return np.abs(values_src - values_dst) + distances 35 36 g, n = pixel_graph( 37 image, mask=mask, connectivity=2, edge_function=edge_func 38 ) 39 s2 = np.sqrt(2) 40 np.testing.assert_allclose(g[0, 1], np.abs(image[0, 0] - image[1, 1]) + s2) 41 np.testing.assert_allclose(g[1, 2], np.abs(image[1, 1] - image[1, 2]) + 1) 42 np.testing.assert_array_equal(n, [0, 4, 5, 7]) 43 44 45def test_default_edge_func(): 46 g, n = pixel_graph(image, spacing=np.array([0.78, 0.78])) 47 num_edges = len(g.data) // 2 # each edge appears in both directions 48 assert num_edges == 12 # lattice in a (3, 3) grid 49 np.testing.assert_almost_equal( 50 g[0, 1], 0.78 * np.abs(image[0, 0] - image[0, 1]) 51 ) 52 np.testing.assert_array_equal(n, np.arange(image.size)) 53 54 55if __name__ == '__main__': 56 test_edge_function() 57