Coverage for src/abm_initialization_collection/image/create_voronoi_image.py: 100%

1from math import floor 

2 

3import numpy as np 

4from bioio import BioImage 

5from scipy.ndimage import binary_dilation, binary_fill_holes, distance_transform_edt 

6 

7 

8def create_voronoi_image(image: BioImage, channel: int, iterations: int, height: int) -> np.ndarray: 

9 """ 

10 Apply Voronoi tessellation to image. 

11 

12 Parameters 

13 ---------- 

14 image 

15 Segmentation image. 

16 channel 

17 Image channel. 

18 iterations 

19 Number of boundary estimation steps. 

20 height 

21 Target height in voxels. 

22 

23 Returns 

24 ------- 

25 : 

26 Voronoi tessellation. 

27 """ 

28 

29 array = image.get_image_data("ZYX", T=0, C=channel) 

30 

31 # Create artificial boundary for voronoi. 

32 mask = create_boundary_mask(array, iterations) 

33 lower_bound, upper_bound = get_mask_bounds(array, height) 

34 mask_id = np.iinfo(array.dtype).max 

35 array[mask == 0] = mask_id 

36 mask[:lower_bound, :, :] = 0 

37 mask[upper_bound:, :, :] = 0 

38 

39 # Calculate voronoi on bounded array. 

40 zslice, yslice, xslice = get_array_slices(mask) 

41 voronoi = calculate_voronoi_array(array[zslice, yslice, xslice]) 

42 

43 # Remove masking ids. 

44 array[zslice, yslice, xslice] = voronoi 

45 array[mask == 0] = 0 

46 array[array == mask_id] = 0 

47 

48 return array 

49 

50 

51def create_boundary_mask(array: np.ndarray, iterations: int) -> np.ndarray: 

52 """ 

53 Creates filled boundary mask around regions in array. 

54 

55 Parameters 

56 ---------- 

57 array 

58 Image array. 

59 iterations 

60 Number of boundary estimation steps. 

61 

62 Returns 

63 ------- 

64 : 

65 Boundary mask array. 

66 """ 

67 

68 mask = np.zeros(array.shape, dtype="uint8") 

69 mask[array != 0] = 1 

70 

71 # Expand using binary dilation to create a border. 

72 binary_dilation(mask, output=mask, iterations=iterations) 

73 

74 # Fill holes in the mask in each z slice. 

75 for z in range(array.shape[0]): 

76 binary_fill_holes(mask[z, :, :], output=mask[z, :, :]) 

77 

78 return mask 

79 

80 

81def get_mask_bounds(array: np.ndarray, target_range: int) -> tuple[int, int]: 

82 """ 

83 Calculates the indices of z axis bounds with given target range. 

84 

85 If the current range between z axis bounds (the minimum and maximum 

86 indices in the z axis where there exist non-zero entries) is wider than 

87 the target range, the current bound indices are returned. 

88 

89 Parameters 

90 ---------- 

91 array 

92 Image array. 

93 target_range 

94 Target distance between bounds. 

95 

96 Returns 

97 ------- 

98 : 

99 Lower and upper bound indices. 

100 """ 

101 

102 lower_bound, upper_bound = np.where(np.any(array, axis=(1, 2)))[0][[0, -1]] 

103 current_range = upper_bound - lower_bound + 1 

104 

105 if current_range < target_range: 

106 height_delta = target_range - current_range 

107 lower_offset = floor(height_delta / 2) 

108 upper_offset = height_delta - lower_offset 

109 lower_bound = lower_bound - lower_offset 

110 upper_bound = upper_bound + upper_offset + 1 

111 else: 

112 upper_bound = upper_bound + 1 

113 

114 return (lower_bound, upper_bound) 

115 

116 

117def get_array_slices(array: np.ndarray) -> tuple[slice, slice, slice]: 

118 """ 

119 Calculate bounding box slices around binary array. 

120 

121 Parameters 

122 ---------- 

123 array 

124 Binary array. 

125 

126 Returns 

127 ------- 

128 : 

129 Slices in the z, y, and x directions. 

130 """ 

131 

132 zsize, ysize, xsize = array.shape 

133 

134 zmin, zmax = np.where(np.any(array, axis=(1, 2)))[0][[0, -1]] 

135 ymin, ymax = np.where(np.any(array, axis=(0, 2)))[0][[0, -1]] 

136 xmin, xmax = np.where(np.any(array, axis=(0, 1)))[0][[0, -1]] 

137 

138 zslice = slice(max(zmin - 1, 0), min(zmax + 2, zsize)) 

139 yslice = slice(max(ymin - 1, 0), min(ymax + 2, ysize)) 

140 xslice = slice(max(xmin - 1, 0), min(xmax + 2, xsize)) 

141 

142 slices = (zslice, yslice, xslice) 

143 return slices 

144 

145 

146def calculate_voronoi_array(array: np.ndarray) -> np.ndarray: 

147 """ 

148 Calculates voronoi on image array using distance transform. 

149 

150 Parameters 

151 ---------- 

152 array 

153 Image array. 

154 

155 Returns 

156 ------- 

157 : 

158 Voronoi array. 

159 """ 

160 

161 distances = distance_transform_edt(array == 0, return_distances=False, return_indices=True) 

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

163 

164 coordinates_z = distances[0].flatten() 

165 coordinates_y = distances[1].flatten() 

166 coordinates_x = distances[2].flatten() 

167 voronoi = array[coordinates_z, coordinates_y, coordinates_x].reshape(array.shape) 

168 

169 return voronoi