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

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, RuntimeTestingRegistryBase, ProcessingContract, FunctionMetadata 

18 

19cp = optional_import("cupy") 

20cucim = optional_import("cucim") 

21cucim_skimage = optional_import("cucim.skimage") 

22 

23 

24class CupyRegistry(RuntimeTestingRegistryBase): 

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 def get_module_patterns(self) -> List[str]: 

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

56 return ['cupy', 'cucim'] 

57 

58 def get_display_name(self) -> str: 

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

60 return 'CuPy' 

61 

62 # ===== HOOK IMPLEMENTATIONS ===== 

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

64 try: 

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

66 except Exception as e: 

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

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

69 

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

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

72 

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

74 return image # No preprocessing needed for CuPy 

75 

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

77 # ProcessingContract system handles dimensional behavior - no categorization needed 

78 return result 

79 

80 # ===== LIBRARY-SPECIFIC IMPLEMENTATIONS ===== 

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

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

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

84 

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

86 """Generate tags with GPU tag.""" 

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

88 tags.append("gpu") 

89 return tags 

90 

91 def _stack_2d_results(self, func, test_3d): 

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

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

94 return cp.stack(results) 

95 

96 def _arrays_close(self, arr1, arr2): 

97 """Compare arrays using CuPy.""" 

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

99 

100