Coverage for src/arraybridge/slice_processing.py: 100%

1""" 

2Shared slice-by-slice processing logic for all memory types. 

3 

4This module provides a single implementation of slice-by-slice processing 

5that works for all memory types, eliminating duplication across dtype wrappers. 

6""" 

7 

8from arraybridge.converters import detect_memory_type 

9from arraybridge.stack_utils import stack_slices, unstack_slices 

10 

11 

12def process_slices(image, func, args, kwargs): 

13 """ 

14 Process a 3D array slice-by-slice using the provided function. 

15 

16 This function handles: 

17 - Unstacking 3D arrays into 2D slices 

18 - Processing each slice independently 

19 - Handling functions that return tuples (main output + special outputs) 

20 - Stacking results back into 3D arrays 

21 - Combining special outputs from all slices 

22 

23 Args: 

24 image: 3D array to process 

25 func: Function to apply to each slice 

26 args: Positional arguments to pass to func 

27 kwargs: Keyword arguments to pass to func 

28 

29 Returns: 

30 Processed 3D array, or tuple of (processed_3d_array, special_outputs...) 

31 if func returns tuples 

32 """ 

33 # Detect memory type and use proper OpenHCS utilities 

34 memory_type = detect_memory_type(image) 

35 gpu_id = 0 # Default GPU ID for slice processing 

36 

37 # Unstack 3D array into 2D slices 

38 slices_2d = unstack_slices(image, memory_type, gpu_id) 

39 

40 # Process each slice and handle special outputs 

41 main_outputs = [] 

42 special_outputs_list = [] 

43 

44 for slice_2d in slices_2d: 

45 slice_result = func(slice_2d, *args, **kwargs) 

46 

47 # Check if result is a tuple (indicating special outputs) 

48 if isinstance(slice_result, tuple): 

49 main_outputs.append(slice_result[0]) # First element is main output 

50 special_outputs_list.append(slice_result[1:]) # Rest are special outputs 

51 else: 

52 main_outputs.append(slice_result) # Single output 

53 

54 # Stack main outputs back into 3D array 

55 result = stack_slices(main_outputs, memory_type, gpu_id) 

56 

57 # If we have special outputs, combine them and return tuple 

58 if special_outputs_list: 

59 # Combine special outputs from all slices 

60 combined_special_outputs = [] 

61 num_special_outputs = len(special_outputs_list[0]) 

62 

63 for i in range(num_special_outputs): 

64 # Collect the i-th special output from all slices 

65 special_output_values = [slice_outputs[i] for slice_outputs in special_outputs_list] 

66 combined_special_outputs.append(special_output_values) 

67 

68 # Return tuple: (stacked_main_output, combined_special_output1, # noqa: E501 

69 # combined_special_output2, ...) 

70 return (result, *combined_special_outputs) 

71 

72 return result