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

1from __future__ import annotations 

2 

3from typing import Callable 

4 

5import numpy as np 

6from scipy import ndimage 

7from skimage import measure 

8 

9 

10def get_neighbors_map(array: np.ndarray) -> dict: 

11 """ 

12 Create map of region ids to lists of neighbors. 

13 

14 Each region id is also assigned a group number, where all regions in a given 

15 group are simply connected. 

16 

17 Parameters 

18 ---------- 

19 array 

20 Segmentation array. 

21 

22 Returns 

23 ------- 

24 : 

25 Map of id to group and neighbor ids. 

26 """ 

27 

28 neighbors_map: dict = {cell_id: {} for cell_id in np.unique(array)} 

29 neighbors_map.pop(0, None) 

30 

31 # Create binary mask for array. 

32 mask = np.zeros(array.shape, dtype="int") 

33 mask[array != 0] = 1 

34 

35 # Label connected groups. 

36 labels, groups = measure.label(mask, connectivity=2, return_num=True) 

37 

38 for group in range(1, groups + 1): 

39 group_crop = get_cropped_array(array, group, labels, crop_original=False) 

40 voxel_ids = [i for i in np.unique(group_crop) if i != 0] 

41 

42 # Find neighbors for each voxel id. 

43 for voxel_id in voxel_ids: 

44 voxel_crop = get_cropped_array(group_crop, voxel_id, crop_original=True) 

45 

46 # Apply custom filter to get border locations. 

47 border_mask = ndimage.generic_filter(voxel_crop, _get_voxel_id_filter(voxel_id), size=3) 

48 

49 # Find neighbors overlapping border. 

50 neighbor_list = np.unique(voxel_crop[border_mask == 1]) 

51 neighbor_list = [i for i in neighbor_list if i not in [0, voxel_id]] 

52 neighbors_map[voxel_id] = {"group": group, "neighbors": neighbor_list} 

53 

54 return neighbors_map 

55 

56 

57def _get_voxel_id_filter(voxel_id: int) -> Callable: 

58 """Create filtering lambda for given id.""" 

59 return lambda v: voxel_id in v 

60 

61 

62def get_bounding_box(array: np.ndarray) -> tuple[int, int, int, int, int, int]: 

63 """ 

64 Find bounding box around array. 

65 

66 Bounds are calculated with a one-voxel border, if possible. 

67 

68 Parameters 

69 ---------- 

70 array 

71 Segmentation array. 

72 

73 Returns 

74 ------- 

75 : 

76 The bounding box (xmin, xmax, ymin, ymax, zmin, zmax) indices 

77 """ 

78 

79 x, y, z = array.shape 

80 

81 xbounds = np.any(array, axis=(1, 2)) 

82 ybounds = np.any(array, axis=(0, 2)) 

83 zbounds = np.any(array, axis=(0, 1)) 

84 

85 xmin, xmax = np.where(xbounds)[0][[0, -1]] 

86 ymin, ymax = np.where(ybounds)[0][[0, -1]] 

87 zmin, zmax = np.where(zbounds)[0][[0, -1]] 

88 

89 xmin = max(xmin - 1, 0) 

90 xmax = min(xmax + 2, x) 

91 

92 ymin = max(ymin - 1, 0) 

93 ymax = min(ymax + 2, y) 

94 

95 zmin = max(zmin - 1, 0) 

96 zmax = min(zmax + 2, z) 

97 

98 return xmin, xmax, ymin, ymax, zmin, zmax 

99 

100 

101def get_cropped_array( 

102 array: np.ndarray, label: int, labels: np.ndarray | None = None, *, crop_original: bool 

103) -> np.ndarray: 

104 """ 

105 Crop array around label region. 

106 

107 Parameters 

108 ---------- 

109 array 

110 Array to crop. 

111 label 

112 Region label. 

113 labels 

114 Array of all region labels. 

115 crop_original 

116 True to crop the original array keeping all labels, False otherwise. 

117 

118 Returns 

119 ------- 

120 : 

121 Cropped array. 

122 """ 

123 

124 # Set all voxels not matching label to zero. 

125 array_mask = array.copy() 

126 array_filter = labels if labels is not None else array_mask 

127 array_mask[array_filter != label] = 0 

128 

129 # Crop array to label. 

130 xmin, xmax, ymin, ymax, zmin, zmax = get_bounding_box(array_mask) 

131 

132 if crop_original: 

133 return array[xmin:xmax, ymin:ymax, zmin:zmax] 

134 

135 return array_mask[xmin:xmax, ymin:ymax, zmin:zmax]