Coverage for src/abm_shape_collection/make_voxels_array.py: 100%

17 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2024-09-25 19:34 +0000

1from __future__ import annotations 

2 

3import numpy as np 

4 

5 

6def make_voxels_array( 

7 voxels: list[tuple[int, int, int]], 

8 reference: list[tuple[int, int, int]] | None = None, 

9 scale: int = 1, 

10) -> np.ndarray: 

11 """ 

12 Convert list of voxels to a binary image array. 

13 

14 Parameters 

15 ---------- 

16 voxels 

17 List of (x, y, z) voxels. 

18 reference 

19 List of (x, y, z) reference voxels used to define center and image size. 

20 scale 

21 Scaling factor for image array. 

22 

23 Returns 

24 ------- 

25 : 

26 Binary image array. 

27 """ 

28 

29 # Set reference if not specified. 

30 if reference is None: 

31 reference = voxels 

32 

33 # Center voxels around (0,0,0). 

34 center_x, center_y, center_z = [round(x) for x in np.array(reference).mean(axis=0)] 

35 voxels_centered = [(z - center_z, y - center_y, x - center_x) for x, y, z in voxels] 

36 reference_centered = [(z - center_z, y - center_y, x - center_x) for x, y, z in reference] 

37 

38 # Create empty array. 

39 mins = np.min(reference_centered, axis=0) 

40 maxs = np.max(reference_centered, axis=0) 

41 height, width, length = np.subtract(maxs, mins) + 3 

42 array = np.zeros((height, width, length), dtype=np.uint8) 

43 

44 # Fill in voxel array. 

45 array_offset = [ 

46 (z - mins[0] + 1, y - mins[1] + 1, x - mins[2] + 1) for z, y, x in voxels_centered 

47 ] 

48 array[tuple(np.transpose(array_offset))] = [1] * len(array_offset) 

49 

50 # Scale the array. 

51 if scale > 1: 

52 array = array.repeat(scale, axis=0).repeat(scale, axis=1).repeat(scale, axis=2) 

53 

54 return array