Coverage for src/abm_colony_collection/get_depth_map.py: 100%

1import numpy as np 

2from scipy import ndimage 

3from skimage import measure 

4 

5 

6def get_depth_map(array: np.ndarray, neighbors_map: dict) -> dict: 

7 """ 

8 Get map of id to depth starting from depth = 1 at the edge if the region. 

9 

10 All ids at the edge of the array are assigned a depth of 1. Immediate 

11 neighbors of those edges are assigned a depth of 2, and so on, until all ids 

12 have an assigned depth. 

13 

14 Parameters 

15 ---------- 

16 array 

17 Segmentation array. 

18 neighbors_map 

19 Map of ids to lists of neighbors. 

20 

21 Returns 

22 ------- 

23 : 

24 Map of id to depth from edge. 

25 """ 

26 

27 depth_map = {cell_id: 0 for cell_id in np.unique(array)} 

28 depth_map.pop(0, None) 

29 

30 # Return empty depth map if there are no cell ids in the array 

31 if not depth_map: 

32 return depth_map 

33 

34 edge_ids = find_edge_ids(array) 

35 visited = set(edge_ids) 

36 queue = edge_ids.copy() 

37 

38 while queue: 

39 current_id = queue.pop(0) 

40 

41 current_neighbors = neighbors_map[current_id]["neighbors"] 

42 valid_neighbors = set(current_neighbors) - visited 

43 visited.update(valid_neighbors) 

44 queue = queue + list(valid_neighbors) 

45 

46 for neighbor_id in valid_neighbors: 

47 depth_map[neighbor_id] = depth_map[current_id] + 1 

48 

49 depth_map[current_id] = depth_map[current_id] + 1 

50 

51 return depth_map 

52 

53 

54def find_edge_ids(array: np.ndarray) -> list[int]: 

55 """ 

56 Get ids of regions closest to the edge of the array. 

57 

58 Parameters 

59 ---------- 

60 array 

61 Segmentation array. 

62 

63 Returns 

64 ------- 

65 : 

66 List of edge arrays. 

67 """ 

68 

69 slice_index = np.argmax(np.count_nonzero(array, axis=(1, 2))) 

70 array_slice = array[slice_index, :, :] 

71 

72 # Calculate voronoi from cell shapes. 

73 distances = ndimage.distance_transform_edt( 

74 array_slice == 0, return_distances=False, return_indices=True 

75 ) 

76 distances = distances.astype("uint16", copy=False) 

77 coordinates_y = distances[0].flatten() 

78 coordinates_x = distances[1].flatten() 

79 voronoi = array_slice[coordinates_y, coordinates_x].reshape(array_slice.shape) 

80 

81 # Create border mask. 

82 mask = np.zeros(array_slice.shape, dtype="uint8") 

83 mask[array_slice != 0] = 1 

84 while measure.euler_number(mask) != 1: 

85 mask = ndimage.binary_dilation(mask, iterations=1) 

86 

87 # Filter voronoi by mask to get edge ids. 

88 voronoi[mask == 1] = 0 

89 edge_ids = list(np.unique(voronoi)) 

90 edge_ids.remove(0) 

91 

92 return edge_ids