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

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 numpy as np 

11from typing import Tuple, List 

12 

13from openhcs.constants import MemoryType 

14from openhcs.core.utils import optional_import 

15from .unified_registry import LibraryRegistryBase, RuntimeTestingRegistryBase 

16 

17cp = optional_import("cupy") 

18cucim = optional_import("cucim") 

19cucim_skimage = optional_import("cucim.skimage") 

20 

21 

22class CupyRegistry(RuntimeTestingRegistryBase): 

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

24 

25 # Registry name for auto-registration 

26 _registry_name = 'cupy' 

27 

28 # Library-specific exclusions (uses common ones) 

29 EXCLUSIONS = LibraryRegistryBase.COMMON_EXCLUSIONS 

30 

31 # Modules to scan for functions 

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

33 'feature', 'restoration', 'transform', 'exposure', 

34 'color', 'util'] 

35 

36 # Memory type for this registry 

37 MEMORY_TYPE = MemoryType.CUPY.value 

38 

39 # Float dtype for this registry 

40 FLOAT_DTYPE = cp.float32 

41 

42 def __init__(self): 

43 super().__init__("cupy") 

44 

45 # ===== ESSENTIAL ABC METHODS ===== 

46 def get_library_version(self) -> str: 

47 return cucim.__version__ 

48 

49 def is_library_available(self) -> bool: 

50 return bool(cp) and bool(cucim_skimage) 

51 

52 def get_library_object(self): 

53 return cucim_skimage 

54 

55 def get_module_patterns(self) -> List[str]: 

56 """Get module patterns for CuPy (includes cucim patterns).""" 

57 return ['cupy', 'cucim'] 

58 

59 def get_display_name(self) -> str: 

60 """Get proper display name for CuPy.""" 

61 return 'CuPy' 

62 

63 # ===== HOOK IMPLEMENTATIONS ===== 

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

65 try: 

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

67 except Exception as e: 

68 # If CUDA initialization fails, raise a more descriptive error 

69 raise RuntimeError(f"CUDA initialization failed during CuPy array creation: {e}") from e 

70 

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

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

73 

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

75 return image # No preprocessing needed for CuPy 

76 

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

78 # ProcessingContract system handles dimensional behavior - no categorization needed 

79 return result 

80 

81 # ===== LIBRARY-SPECIFIC IMPLEMENTATIONS ===== 

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

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

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

85 

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

87 """Generate tags with GPU tag.""" 

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

89 tags.append("gpu") 

90 return tags 

91 

92 def _stack_2d_results(self, func, test_3d): 

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

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

95 return cp.stack(results) 

96 

97 def _arrays_close(self, arr1, arr2): 

98 """Compare arrays using CuPy.""" 

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

100 

101