Coverage for openhcs/processing/backends/lib_registry/cupy_registry.py: 71.1%

1""" 

2Clean CuPy Registry Implementation 

3 

4Implements clean abstraction with internal library-specific logic. 

5All CuPy-specific details (GPU handling, CuCIM integration, etc.) 

6are handled internally without leaking into the ABC. 

7""" 

8from __future__ import annotations 

9 

10import inspect 

11import numpy as np 

12from functools import wraps 

13from typing import Tuple, Callable, List, Any, Dict 

14 

15from openhcs.constants import MemoryType 

16from openhcs.core.utils import optional_import 

17from .unified_registry import LibraryRegistryBase, ProcessingContract, FunctionMetadata 

18 

19cp = optional_import("cupy") 

20cucim = optional_import("cucim") 

21cucim_skimage = optional_import("cucim.skimage") 

22 

23 

24class CupyRegistry(LibraryRegistryBase): 

25 """Clean CuPy registry with internal GPU handling logic.""" 

26 

27 # Library-specific exclusions (uses common ones) 

28 EXCLUSIONS = LibraryRegistryBase.COMMON_EXCLUSIONS 

29 

30 # Modules to scan for functions 

31 MODULES_TO_SCAN = ['filters', 'morphology', 'measure', 'segmentation', 

32 'feature', 'restoration', 'transform', 'exposure', 

33 'color', 'util'] 

34 

35 # Memory type for this registry 

36 MEMORY_TYPE = MemoryType.CUPY.value 

37 

38 # Float dtype for this registry 

39 FLOAT_DTYPE = cp.float32 

40 

41 def __init__(self): 

42 super().__init__("cupy") 

43 

44 # ===== ESSENTIAL ABC METHODS ===== 

45 def get_library_version(self) -> str: 

46 return cucim.__version__ 

47 

48 def is_library_available(self) -> bool: 

49 return cp is not None and cucim_skimage is not None 

50 

51 def get_library_object(self): 

52 return cucim_skimage 

53 

54 # ===== HOOK IMPLEMENTATIONS ===== 

55 def _create_array(self, shape: Tuple[int, ...], dtype): 

56 return cp.random.rand(*shape).astype(dtype) 

57 

58 def _check_first_parameter(self, first_param, func_name: str) -> bool: 

59 return first_param.name.lower() in {'image', 'input', 'array', 'img'} 

60 

61 def _preprocess_input(self, image, func_name: str): 

62 return image # No preprocessing needed for CuPy 

63 

64 def _postprocess_output(self, result, original_image, func_name: str): 

65 # ProcessingContract system handles dimensional behavior - no categorization needed 

66 return result 

67 

68 # ===== LIBRARY-SPECIFIC IMPLEMENTATIONS ===== 

69 def _generate_function_name(self, name: str, module_name: str) -> str: 

70 """Generate function name - original for filters, prefixed for others.""" 

71 return name if module_name == 'filters' else f"{module_name}_{name}" 

72 

73 def _generate_tags(self, func_name: str) -> List[str]: 

74 """Generate tags with GPU tag.""" 

75 tags = func_name.lower().replace("_", " ").split() 

76 tags.append("gpu") 

77 return tags 

78 

79 def _stack_2d_results(self, func, test_3d): 

80 """Stack 2D results using CuPy.""" 

81 results = [func(test_3d[z]) for z in range(test_3d.shape[0])] 

82 return cp.stack(results) 

83 

84 def _arrays_close(self, arr1, arr2): 

85 """Compare arrays using CuPy.""" 

86 return np.allclose(arr1.get(), arr2.get(), rtol=1e-5, atol=1e-8) 

87 

88 def _expand_2d_to_3d(self, array_2d): 

89 """Expand 2D array to 3D using CuPy expansion.""" 

90 return array_2d[None, ...]