Coverage for openhcs/processing/backends/processors/cupy_processor.py: 14.3%

1""" 

2CuPy Image Processor Implementation 

3 

4This module implements the ImageProcessorInterface using CuPy as the backend. 

5It leverages GPU acceleration for image processing operations. 

6 

7Doctrinal Clauses: 

8- Clause 3 — Declarative Primacy: All functions are pure and stateless 

9- Clause 88 — No Inferred Capabilities: Explicit CuPy dependency 

10- Clause 106-A — Declared Memory Types: All methods specify CuPy arrays 

11""" 

12from __future__ import annotations 

13 

14import logging 

15from typing import Any, List, Optional, Tuple 

16 

17from openhcs.core.memory.decorators import cupy as cupy_func 

18from openhcs.core.utils import optional_import 

19 

20# Import CuPy as an optional dependency 

21cp = optional_import("cupy") 

22ndimage = None 

23if cp is not None: 23 ↛ 29line 23 didn't jump to line 29 because the condition on line 23 was always true

24 cupyx_scipy = optional_import("cupyx.scipy") 

25 if cupyx_scipy is not None: 25 ↛ 29line 25 didn't jump to line 29 because the condition on line 25 was always true

26 ndimage = cupyx_scipy.ndimage 

27 

28# Import CuCIM for edge detection 

29cucim_filters = optional_import("cucim.skimage.filters") 

30 

31logger = logging.getLogger(__name__) 

32 

33 

34@cupy_func 

35def create_linear_weight_mask(height: int, width: int, margin_ratio: float = 0.1) -> "cp.ndarray": 

36 """ 

37 Create a 2D weight mask that linearly ramps from 0 at the edges to 1 in the center. 

38 

39 Args: 

40 height: Height of the mask 

41 width: Width of the mask 

42 margin_ratio: Ratio of the margin to the image size 

43 

44 Returns: 

45 2D CuPy weight mask of shape (height, width) 

46 """ 

47 # The compiler will ensure this function is only called when CuPy is available 

48 # No need to check for CuPy availability here 

49 

50 margin_y = int(cp.floor(height * margin_ratio)) 

51 margin_x = int(cp.floor(width * margin_ratio)) 

52 

53 weight_y = cp.ones(height, dtype=cp.float32) 

54 if margin_y > 0: 

55 ramp_top = cp.linspace(0, 1, margin_y, endpoint=False) 

56 ramp_bottom = cp.linspace(1, 0, margin_y, endpoint=False) 

57 weight_y[:margin_y] = ramp_top 

58 weight_y[-margin_y:] = ramp_bottom 

59 

60 weight_x = cp.ones(width, dtype=cp.float32) 

61 if margin_x > 0: 

62 ramp_left = cp.linspace(0, 1, margin_x, endpoint=False) 

63 ramp_right = cp.linspace(1, 0, margin_x, endpoint=False) 

64 weight_x[:margin_x] = ramp_left 

65 weight_x[-margin_x:] = ramp_right 

66 

67 # Create 2D weight mask 

68 weight_mask = cp.outer(weight_y, weight_x) 

69 

70 return weight_mask 

71 

72 

73def _validate_3d_array(array: Any, name: str = "input") -> None: 

74 """ 

75 Validate that the input is a 3D CuPy array. 

76 

77 Args: 

78 array: Array to validate 

79 name: Name of the array for error messages 

80 

81 Raises: 

82 TypeError: If the array is not a CuPy array 

83 ValueError: If the array is not 3D 

84 ImportError: If CuPy is not available 

85 """ 

86 # The compiler will ensure this function is only called when CuPy is available 

87 # No need to check for CuPy availability here 

88 

89 if not isinstance(array, cp.ndarray): 

90 raise TypeError(f"{name} must be a CuPy array, got {type(array)}. " 

91 f"No automatic conversion is performed to maintain explicit contracts.") 

92 

93 if array.ndim != 3: 

94 raise ValueError(f"{name} must be a 3D array, got {array.ndim}D") 

95 

96@cupy_func 

97def sharpen(image: "cp.ndarray", radius: float = 1.0, amount: float = 1.0) -> "cp.ndarray": 

98 """ 

99 Sharpen a 3D image using unsharp masking - GPU PARALLELIZED. 

100 

101 This applies sharpening to each Z-slice independently using vectorized operations 

102 for maximum GPU utilization. Normalization and rescaling are fully parallelized. 

103 

104 Args: 

105 image: 3D CuPy array of shape (Z, Y, X) 

106 radius: Radius of Gaussian blur 

107 amount: Sharpening strength 

108 

109 Returns: 

110 Sharpened 3D CuPy array of shape (Z, Y, X) 

111 """ 

112 _validate_3d_array(image) 

113 

114 # Store original dtype 

115 dtype = image.dtype 

116 

117 # Convert to float32 for processing 

118 image_float = image.astype(cp.float32) 

119 

120 # Vectorized per-slice normalization - GPU PARALLELIZED 

121 # Get max value per slice: shape (Z, 1, 1) for broadcasting 

122 max_per_slice = cp.max(image_float, axis=(1, 2), keepdims=True) 

123 image_norm = image_float / max_per_slice 

124 

125 # Apply 3D Gaussian blur with sigma_z=0 for slice-wise processing - GPU PARALLELIZED 

126 # This processes all slices simultaneously while keeping Z-slices independent 

127 blurred = ndimage.gaussian_filter(image_norm, sigma=(0, radius, radius)) 

128 

129 # Apply unsharp mask: original + amount * (original - blurred) - GPU PARALLELIZED 

130 sharpened = image_norm + amount * (image_norm - blurred) 

131 

132 # Clip to valid range - GPU PARALLELIZED 

133 sharpened = cp.clip(sharpened, 0, 1.0) 

134 

135 # Vectorized rescaling back to original range - GPU PARALLELIZED 

136 min_per_slice = cp.min(sharpened, axis=(1, 2), keepdims=True) 

137 max_per_slice = cp.max(sharpened, axis=(1, 2), keepdims=True) 

138 

139 # Avoid division by zero using broadcasting 

140 range_per_slice = max_per_slice - min_per_slice 

141 # Only rescale slices where max > min 

142 valid_range = range_per_slice > 0 

143 sharpened_rescaled = cp.where( 

144 valid_range, 

145 (sharpened - min_per_slice) * 65535 / range_per_slice, 

146 sharpened * 65535 

147 ) 

148 

149 # Convert back to original dtype 

150 return sharpened_rescaled.astype(dtype) 

151 

152@cupy_func 

153def percentile_normalize( 

154 image: "cp.ndarray", 

155 low_percentile: float = 1.0, 

156 high_percentile: float = 99.0, 

157 target_min: float = 0.0, 

158 target_max: float = 65535.0 

159) -> "cp.ndarray": 

160 """ 

161 Normalize a 3D image using percentile-based contrast stretching. 

162 

163 This applies normalization to each Z-slice independently using slice-by-slice 

164 processing for algorithmic reasons (each slice needs different percentile values). 

165 Each slice operation is GPU parallelized across Y,X dimensions. 

166 

167 Args: 

168 image: 3D CuPy array of shape (Z, Y, X) 

169 low_percentile: Lower percentile (0-100) 

170 high_percentile: Upper percentile (0-100) 

171 target_min: Target minimum value 

172 target_max: Target maximum value 

173 

174 Returns: 

175 Normalized 3D CuPy array of shape (Z, Y, X) 

176 """ 

177 _validate_3d_array(image) 

178 

179 # Process each Z-slice independently - MATCH NUMPY EXACTLY 

180 # NOTE: For loop is ALGORITHMICALLY NECESSARY here because each slice needs 

181 # different percentile values. Cannot be vectorized without changing the algorithm. 

182 result = cp.zeros_like(image, dtype=cp.float32) # Use float32 like NumPy 

183 

184 for z in range(image.shape[0]): 

185 # Get percentile values for this slice - MATCH NUMPY EXACTLY 

186 p_low, p_high = cp.percentile(image[z], (low_percentile, high_percentile)) 

187 

188 # Avoid division by zero - MATCH NUMPY EXACTLY 

189 if p_high == p_low: 

190 result[z] = cp.ones_like(image[z]) * target_min 

191 continue 

192 

193 # Clip and normalize to target range - MATCH NUMPY EXACTLY 

194 clipped = cp.clip(image[z], p_low, p_high) 

195 normalized = (clipped - p_low) * (target_max - target_min) / (p_high - p_low) + target_min 

196 result[z] = normalized 

197 

198 # Convert to uint16 - MATCH NUMPY EXACTLY 

199 return result.astype(cp.uint16) 

200 

201@cupy_func 

202def stack_percentile_normalize( 

203 stack: "cp.ndarray", 

204 low_percentile: float = 1.0, 

205 high_percentile: float = 99.0, 

206 target_min: float = 0.0, 

207 target_max: float = 65535.0 

208) -> "cp.ndarray": 

209 """ 

210 Normalize a stack using global percentile-based contrast stretching. 

211 

212 This ensures consistent normalization across all Z-slices by computing 

213 global percentiles across the entire stack. 

214 

215 Args: 

216 stack: 3D CuPy array of shape (Z, Y, X) 

217 low_percentile: Lower percentile (0-100) 

218 high_percentile: Upper percentile (0-100) 

219 target_min: Target minimum value 

220 target_max: Target maximum value 

221 

222 Returns: 

223 Normalized 3D CuPy array of shape (Z, Y, X) 

224 """ 

225 _validate_3d_array(stack) 

226 

227 # Calculate global percentiles across the entire stack 

228 p_low = cp.percentile(stack, low_percentile) 

229 p_high = cp.percentile(stack, high_percentile) 

230 

231 # Avoid division by zero 

232 if p_high == p_low: 

233 return cp.ones_like(stack) * target_min 

234 

235 # Clip and normalize to target range (match NumPy implementation exactly) 

236 clipped = cp.clip(stack, p_low, p_high) 

237 normalized = (clipped - p_low) * (target_max - target_min) / (p_high - p_low) + target_min 

238 normalized = normalized.astype(cp.uint16) 

239 

240 return normalized 

241 

242@cupy_func 

243def create_composite( 

244 stack: "cp.ndarray", weights: Optional[List[float]] = None 

245) -> "cp.ndarray": 

246 """ 

247 Create a composite image from a 3D stack where each slice is a channel. 

248 

249 Args: 

250 stack: 3D CuPy array of shape (N, Y, X) where N is number of channel slices 

251 weights: List of weights for each slice. If None, equal weights are used. 

252 

253 Returns: 

254 Composite 3D CuPy array of shape (1, Y, X) 

255 """ 

256 # 🔄 MEMORY CONVERSION LOGGING: Log what we actually received 

257 import logging 

258 logger = logging.getLogger(__name__) 

259 logger.info(f"🔄 CREATE_COMPOSITE: Called with stack type: {type(stack)}, shape: {getattr(stack, 'shape', 'no shape')}") 

260 logger.info(f"🔄 CREATE_COMPOSITE: weights parameter - type: {type(weights)}, value: {weights}") 

261 

262 # Validate input is 3D array 

263 if not hasattr(stack, 'shape') or len(stack.shape) != 3: 

264 raise TypeError(f"stack must be a 3D CuPy array, got shape: {getattr(stack, 'shape', 'no shape')}") 

265 

266 n_slices, height, width = stack.shape 

267 

268 # Default weights if none provided 

269 if weights is None: 

270 # Equal weights for all slices 

271 weights = [1.0 / n_slices] * n_slices 

272 logger.info(f"🔄 CREATE_COMPOSITE: Using default equal weights: {weights}") 

273 elif isinstance(weights, (list, tuple)): 

274 # Convert tuple to list if needed 

275 weights = list(weights) 

276 logger.info(f"🔄 CREATE_COMPOSITE: Using provided weights: {weights}") 

277 if len(weights) != n_slices: 

278 raise ValueError(f"Number of weights ({len(weights)}) must match number of slices ({n_slices})") 

279 else: 

280 # Log the problematic type and value for debugging 

281 logger.error(f"🔄 CREATE_COMPOSITE: Invalid weights type - expected list/tuple/None, got {type(weights)}: {weights}") 

282 raise TypeError(f"weights must be a list of values or None, got {type(weights)}: {weights}") 

283 

284 # Normalize weights to sum to 1 

285 weight_sum = sum(weights) 

286 if weight_sum == 0: 

287 raise ValueError("Sum of weights cannot be zero") 

288 normalized_weights = [w / weight_sum for w in weights] 

289 

290 # Convert weights to CuPy array for efficient computation 

291 # CRITICAL: Use float32 for weights to preserve fractional values, not stack.dtype 

292 weights_array = cp.array(normalized_weights, dtype=cp.float32) 

293 

294 # Reshape weights for broadcasting: (N, 1, 1) to multiply with (N, Y, X) 

295 weights_array = weights_array.reshape(n_slices, 1, 1) 

296 

297 # Create composite by weighted sum along the first axis 

298 # Convert stack to float32 for computation to avoid precision loss 

299 stack_float = stack.astype(cp.float32) 

300 weighted_stack = stack_float * weights_array 

301 composite_slice = cp.sum(weighted_stack, axis=0, keepdims=True) # Keep as (1, Y, X) 

302 

303 # Convert back to original dtype 

304 composite_slice = composite_slice.astype(stack.dtype) 

305 

306 return composite_slice 

307 

308@cupy_func 

309def apply_mask(image: "cp.ndarray", mask: "cp.ndarray") -> "cp.ndarray": 

310 """ 

311 Apply a mask to a 3D image. 

312 

313 This applies the mask to each Z-slice independently if mask is 2D, 

314 or applies the 3D mask directly if mask is 3D. 

315 

316 Args: 

317 image: 3D CuPy array of shape (Z, Y, X) 

318 mask: 3D CuPy array of shape (Z, Y, X) or 2D CuPy array of shape (Y, X) 

319 

320 Returns: 

321 Masked 3D CuPy array of shape (Z, Y, X) 

322 """ 

323 _validate_3d_array(image) 

324 

325 # Handle 2D mask (apply to each Z-slice) 

326 if isinstance(mask, cp.ndarray) and mask.ndim == 2: 

327 if mask.shape != image.shape[1:]: 

328 raise ValueError( 

329 f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" 

330 ) 

331 

332 # Apply 2D mask to all Z-slices simultaneously using broadcasting - GPU PARALLELIZED 

333 # Broadcasting mask from (Y, X) to (1, Y, X) allows vectorized operation across all slices 

334 mask_3d = mask[None, :, :] # Shape: (1, Y, X) 

335 result = image.astype(cp.float32) * mask_3d.astype(cp.float32) 

336 

337 return result.astype(image.dtype) 

338 

339 # Handle 3D mask 

340 if isinstance(mask, cp.ndarray) and mask.ndim == 3: 

341 if mask.shape != image.shape: 

342 raise ValueError( 

343 f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" 

344 ) 

345 

346 # Apply 3D mask directly - MATCH NUMPY EXACTLY 

347 masked = image.astype(cp.float32) * mask.astype(cp.float32) 

348 return masked.astype(image.dtype) 

349 

350 # If we get here, the mask is neither 2D nor 3D CuPy array 

351 raise TypeError(f"mask must be a 2D or 3D CuPy array, got {type(mask)}") 

352 

353@cupy_func 

354def create_weight_mask(shape: Tuple[int, int], margin_ratio: float = 0.1) -> "cp.ndarray": 

355 """ 

356 Create a weight mask for blending images. 

357 

358 Args: 

359 shape: Shape of the mask (height, width) 

360 margin_ratio: Ratio of image size to use as margin 

361 

362 Returns: 

363 2D CuPy weight mask of shape (Y, X) 

364 """ 

365 if not isinstance(shape, tuple) or len(shape) != 2: 

366 raise TypeError("shape must be a tuple of (height, width)") 

367 

368 height, width = shape 

369 return create_linear_weight_mask(height, width, margin_ratio) 

370 

371@cupy_func 

372def max_projection(stack: "cp.ndarray") -> "cp.ndarray": 

373 """ 

374 Create a maximum intensity projection from a Z-stack. 

375 

376 Args: 

377 stack: 3D CuPy array of shape (Z, Y, X) 

378 

379 Returns: 

380 3D CuPy array of shape (1, Y, X) 

381 """ 

382 _validate_3d_array(stack) 

383 

384 # Create max projection 

385 projection_2d = cp.max(stack, axis=0) 

386 return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) 

387 

388@cupy_func 

389def mean_projection(stack: "cp.ndarray") -> "cp.ndarray": 

390 """ 

391 Create a mean intensity projection from a Z-stack. 

392 

393 Args: 

394 stack: 3D CuPy array of shape (Z, Y, X) 

395 

396 Returns: 

397 3D CuPy array of shape (1, Y, X) 

398 """ 

399 _validate_3d_array(stack) 

400 

401 # Create mean projection 

402 projection_2d = cp.mean(stack, axis=0).astype(stack.dtype) 

403 return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) 

404 

405@cupy_func 

406def spatial_bin_2d( 

407 stack: "cp.ndarray", 

408 bin_size: int = 2, 

409 method: str = "mean" 

410) -> "cp.ndarray": 

411 """ 

412 Apply 2D spatial binning to each slice in the stack - GPU accelerated. 

413 

414 Reduces spatial resolution by combining neighboring pixels in 2D blocks. 

415 Each slice is processed independently using efficient GPU operations. 

416 

417 Args: 

418 stack: 3D CuPy array of shape (Z, Y, X) 

419 bin_size: Size of the square binning kernel (e.g., 2 = 2x2 binning) 

420 method: Binning method - "mean", "sum", "max", or "min" 

421 

422 Returns: 

423 Binned 3D CuPy array of shape (Z, Y//bin_size, X//bin_size) 

424 """ 

425 _validate_3d_array(stack) 

426 

427 if bin_size <= 0: 

428 raise ValueError("bin_size must be positive") 

429 if method not in ["mean", "sum", "max", "min"]: 

430 raise ValueError("method must be one of: mean, sum, max, min") 

431 

432 z_slices, height, width = stack.shape 

433 

434 # Calculate output dimensions 

435 new_height = height // bin_size 

436 new_width = width // bin_size 

437 

438 if new_height == 0 or new_width == 0: 

439 raise ValueError(f"bin_size {bin_size} is too large for image dimensions {height}x{width}") 

440 

441 # Crop to make dimensions divisible by bin_size 

442 crop_height = new_height * bin_size 

443 crop_width = new_width * bin_size 

444 cropped_stack = stack[:, :crop_height, :crop_width] 

445 

446 # Reshape for binning: (Z, new_height, bin_size, new_width, bin_size) 

447 reshaped = cropped_stack.reshape(z_slices, new_height, bin_size, new_width, bin_size) 

448 

449 # Apply binning operation using CuPy functions 

450 if method == "mean": 

451 result = cp.mean(reshaped, axis=(2, 4)) 

452 elif method == "sum": 

453 result = cp.sum(reshaped, axis=(2, 4)) 

454 elif method == "max": 

455 result = cp.max(reshaped, axis=(2, 4)) 

456 elif method == "min": 

457 result = cp.min(reshaped, axis=(2, 4)) 

458 

459 return result.astype(stack.dtype) 

460 

461@cupy_func 

462def spatial_bin_3d( 

463 stack: "cp.ndarray", 

464 bin_size: int = 2, 

465 method: str = "mean" 

466) -> "cp.ndarray": 

467 """ 

468 Apply 3D spatial binning to the entire stack - GPU accelerated. 

469 

470 Reduces spatial resolution by combining neighboring voxels in 3D blocks 

471 using efficient GPU operations. 

472 

473 Args: 

474 stack: 3D CuPy array of shape (Z, Y, X) 

475 bin_size: Size of the cubic binning kernel (e.g., 2 = 2x2x2 binning) 

476 method: Binning method - "mean", "sum", "max", or "min" 

477 

478 Returns: 

479 Binned 3D CuPy array of shape (Z//bin_size, Y//bin_size, X//bin_size) 

480 """ 

481 _validate_3d_array(stack) 

482 

483 if bin_size <= 0: 

484 raise ValueError("bin_size must be positive") 

485 if method not in ["mean", "sum", "max", "min"]: 

486 raise ValueError("method must be one of: mean, sum, max, min") 

487 

488 depth, height, width = stack.shape 

489 

490 # Calculate output dimensions 

491 new_depth = depth // bin_size 

492 new_height = height // bin_size 

493 new_width = width // bin_size 

494 

495 if new_depth == 0 or new_height == 0 or new_width == 0: 

496 raise ValueError(f"bin_size {bin_size} is too large for stack dimensions {depth}x{height}x{width}") 

497 

498 # Crop to make dimensions divisible by bin_size 

499 crop_depth = new_depth * bin_size 

500 crop_height = new_height * bin_size 

501 crop_width = new_width * bin_size 

502 cropped_stack = stack[:crop_depth, :crop_height, :crop_width] 

503 

504 # Reshape for 3D binning: (new_depth, bin_size, new_height, bin_size, new_width, bin_size) 

505 reshaped = cropped_stack.reshape(new_depth, bin_size, new_height, bin_size, new_width, bin_size) 

506 

507 # Apply binning operation across the three bin_size dimensions using CuPy functions 

508 if method == "mean": 

509 result = cp.mean(reshaped, axis=(1, 3, 5)) 

510 elif method == "sum": 

511 result = cp.sum(reshaped, axis=(1, 3, 5)) 

512 elif method == "max": 

513 result = cp.max(reshaped, axis=(1, 3, 5)) 

514 elif method == "min": 

515 result = cp.min(reshaped, axis=(1, 3, 5)) 

516 

517 return result.astype(stack.dtype) 

518 

519@cupy_func 

520def stack_equalize_histogram( 

521 stack: "cp.ndarray", 

522 bins: int = 65536, 

523 range_min: float = 0.0, 

524 range_max: float = 65535.0 

525) -> "cp.ndarray": 

526 """ 

527 Apply histogram equalization to an entire stack. 

528 

529 This ensures consistent contrast enhancement across all Z-slices by 

530 computing a global histogram across the entire stack. 

531 

532 Args: 

533 stack: 3D CuPy array of shape (Z, Y, X) 

534 bins: Number of bins for histogram computation 

535 range_min: Minimum value for histogram range 

536 range_max: Maximum value for histogram range 

537 

538 Returns: 

539 Equalized 3D CuPy array of shape (Z, Y, X) 

540 """ 

541 _validate_3d_array(stack) 

542 

543 # MATCH NUMPY EXACTLY - Flatten the entire stack to compute the global histogram 

544 flat_stack = stack.flatten() 

545 

546 # Calculate the histogram and cumulative distribution function (CDF) - MATCH NUMPY EXACTLY 

547 hist, bin_edges = cp.histogram(flat_stack, bins=bins, range=(range_min, range_max)) 

548 cdf = hist.cumsum() 

549 

550 # Normalize the CDF to the range [0, 65535] - MATCH NUMPY EXACTLY 

551 # Avoid division by zero 

552 if cdf[-1] > 0: 

553 cdf = 65535 * cdf / cdf[-1] 

554 

555 # Use linear interpolation to map input values to equalized values - MATCH NUMPY EXACTLY 

556 equalized_stack = cp.interp(stack.flatten(), bin_edges[:-1], cdf).reshape(stack.shape) 

557 

558 # Convert to uint16 - MATCH NUMPY EXACTLY 

559 return equalized_stack.astype(cp.uint16) 

560 

561@cupy_func 

562def create_projection( 

563 stack: "cp.ndarray", method: str = "max_projection" 

564) -> "cp.ndarray": 

565 """ 

566 Create a projection from a stack using the specified method. 

567 

568 Args: 

569 stack: 3D CuPy array of shape (Z, Y, X) 

570 method: Projection method (max_projection, mean_projection) 

571 

572 Returns: 

573 3D CuPy array of shape (1, Y, X) 

574 """ 

575 _validate_3d_array(stack) 

576 

577 if method == "max_projection": 

578 return max_projection(stack) 

579 

580 if method == "mean_projection": 

581 return mean_projection(stack) 

582 

583 # FAIL FAST: No fallback projection methods 

584 raise ValueError(f"Unknown projection method: {method}. Valid methods: max_projection, mean_projection") 

585 

586 

587@cupy_func 

588def crop( 

589 input_image: "cp.ndarray", 

590 start_x: int = 0, 

591 start_y: int = 0, 

592 start_z: int = 0, 

593 width: int = 1, 

594 height: int = 1, 

595 depth: int = 1 

596) -> "cp.ndarray": 

597 """ 

598 Crop a given substack out of a given image stack. 

599 

600 Equivalent to pyclesperanto.crop() but using CuPy operations. 

601 

602 Parameters 

603 ---------- 

604 input_image: cp.ndarray 

605 Input 3D image to process of shape (Z, Y, X) 

606 start_x: int (= 0) 

607 Starting index coordinate x 

608 start_y: int (= 0) 

609 Starting index coordinate y 

610 start_z: int (= 0) 

611 Starting index coordinate z 

612 width: int (= 1) 

613 Width size of the region to crop 

614 height: int (= 1) 

615 Height size of the region to crop 

616 depth: int (= 1) 

617 Depth size of the region to crop 

618 

619 Returns 

620 ------- 

621 cp.ndarray 

622 Cropped 3D array of shape (depth, height, width) 

623 """ 

624 _validate_3d_array(input_image) 

625 

626 # Validate crop parameters 

627 if width <= 0 or height <= 0 or depth <= 0: 

628 raise ValueError(f"Crop dimensions must be positive: width={width}, height={height}, depth={depth}") 

629 

630 if start_x < 0 or start_y < 0 or start_z < 0: 

631 raise ValueError(f"Start coordinates must be non-negative: start_x={start_x}, start_y={start_y}, start_z={start_z}") 

632 

633 # Get input dimensions 

634 input_depth, input_height, input_width = input_image.shape 

635 

636 # Calculate end coordinates 

637 end_x = start_x + width 

638 end_y = start_y + height 

639 end_z = start_z + depth 

640 

641 # Validate bounds 

642 if end_x > input_width or end_y > input_height or end_z > input_depth: 

643 raise ValueError( 

644 f"Crop region extends beyond image bounds. " 

645 f"Image shape: {input_image.shape}, " 

646 f"Crop region: ({start_z}:{end_z}, {start_y}:{end_y}, {start_x}:{end_x})" 

647 ) 

648 

649 # Perform the crop using CuPy slicing 

650 cropped = input_image[start_z:end_z, start_y:end_y, start_x:end_x] 

651 

652 return cropped 

653 

654def _create_disk_cupy(radius: int) -> "cp.ndarray": 

655 """Create a disk structuring element using CuPy - MATCH NUMPY EXACTLY""" 

656 y, x = cp.ogrid[-radius:radius+1, -radius:radius+1] 

657 mask = x*x + y*y <= radius*radius 

658 return mask.astype(cp.uint8) 

659 

660def _resize_cupy_better_match(image: "cp.ndarray", output_shape: tuple, anti_aliasing: bool = True, preserve_range: bool = True) -> "cp.ndarray": 

661 """ 

662 Resize image using CuPy to better match scikit-image's transform.resize behavior. 

663 

664 Key differences from original: 

665 1. Better anti-aliasing sigma calculation 

666 2. Proper preserve_range handling without rescaling 

667 3. More accurate zoom parameters 

668 """ 

669 from cupyx.scipy import ndimage as cupy_ndimage 

670 

671 # Calculate zoom factors 

672 zoom_factors = [output_shape[i] / image.shape[i] for i in range(len(output_shape))] 

673 

674 # Convert to float32 for processing 

675 image_float = image.astype(cp.float32) 

676 

677 # Apply anti-aliasing if downsampling 

678 if anti_aliasing and any(z < 1.0 for z in zoom_factors): 

679 # Use scikit-image's sigma calculation: sigma = (1/zoom - 1) / 2 

680 # But ensure minimum sigma and handle edge cases 

681 sigma = [] 

682 for z in zoom_factors: 

683 if z < 1.0: 

684 s = (1.0/z - 1.0) / 2.0 

685 sigma.append(max(s, 0.5)) # Minimum sigma for stability 

686 else: 

687 sigma.append(0.0) 

688 

689 # Apply Gaussian smoothing before downsampling 

690 if any(s > 0 for s in sigma): 

691 image_float = cupy_ndimage.gaussian_filter(image_float, sigma) 

692 

693 # Perform zoom with bilinear interpolation (order=1) 

694 resized = cupy_ndimage.zoom(image_float, zoom_factors, order=1) 

695 

696 # Handle preserve_range properly - don't rescale, just maintain dtype range 

697 if preserve_range: 

698 # Clip to valid range for the original dtype 

699 if image.dtype == cp.uint16: 

700 resized = cp.clip(resized, 0, 65535) 

701 elif image.dtype == cp.uint8: 

702 resized = cp.clip(resized, 0, 255) 

703 # For other dtypes, keep as-is 

704 

705 return resized 

706 

707@cupy_func 

708def tophat( 

709 image: "cp.ndarray", 

710 selem_radius: int = 50, 

711 downsample_factor: int = 4, 

712 downsample_anti_aliasing: bool = True, 

713 upsample_anti_aliasing: bool = False 

714) -> "cp.ndarray": 

715 """ 

716 Apply white top-hat filter to a 3D image for background removal. 

717 

718 This applies the filter to each Z-slice independently using slice-by-slice 

719 processing for algorithmic reasons (complex multi-step processing with 

720 slice-specific intermediate results). Each slice operation is GPU parallelized. 

721 

722 Args: 

723 image: 3D CuPy array of shape (Z, Y, X) 

724 selem_radius: Radius of the structuring element disk 

725 downsample_factor: Factor by which to downsample the image for processing 

726 downsample_anti_aliasing: Whether to use anti-aliasing when downsampling 

727 upsample_anti_aliasing: Whether to use anti-aliasing when upsampling 

728 

729 Returns: 

730 Filtered 3D CuPy array of shape (Z, Y, X) 

731 """ 

732 _validate_3d_array(image) 

733 

734 # Process each Z-slice independently - IMPROVED MATCH TO NUMPY 

735 # NOTE: For loop is ALGORITHMICALLY NECESSARY here due to complex multi-step 

736 # processing (downsample, morphology, upsample, subtract) with slice-specific 

737 # intermediate results. Vectorization would require significant algorithm restructuring. 

738 result = cp.zeros_like(image) 

739 

740 for z in range(image.shape[0]): 

741 # Store original data type - MATCH NUMPY EXACTLY 

742 input_dtype = image[z].dtype 

743 

744 # 1) Downsample - IMPROVED MATCH TO NUMPY 

745 target_shape = (image[z].shape[0]//downsample_factor, image[z].shape[1]//downsample_factor) 

746 image_small = _resize_cupy_better_match( 

747 image[z], 

748 target_shape, 

749 anti_aliasing=downsample_anti_aliasing, 

750 preserve_range=True 

751 ) 

752 

753 # 2) Build structuring element for the smaller image - MATCH NUMPY EXACTLY 

754 selem_small = _create_disk_cupy(selem_radius // downsample_factor) 

755 

756 # 3) White top-hat on the smaller image - MATCH NUMPY EXACTLY 

757 tophat_small = ndimage.white_tophat(image_small, structure=selem_small) 

758 

759 # 4) Upscale background to original size - IMPROVED MATCH TO NUMPY 

760 background_small = image_small - tophat_small 

761 background_large = _resize_cupy_better_match( 

762 background_small, 

763 image[z].shape, 

764 anti_aliasing=upsample_anti_aliasing, 

765 preserve_range=True 

766 ) 

767 

768 # 5) Subtract background and clip negative values - MATCH NUMPY EXACTLY 

769 slice_result = cp.maximum(image[z] - background_large, 0) 

770 

771 # 6) Convert back to original data type - MATCH NUMPY EXACTLY 

772 result[z] = slice_result.astype(input_dtype) 

773 

774 return result 

775 

776# Lazy initialization of ElementwiseKernel to avoid import errors 

777_sobel_2d_parallel = None 

778 

779def _get_sobel_2d_kernel(): 

780 """Get or create the Sobel 2D parallel kernel.""" 

781 global _sobel_2d_parallel 

782 if _sobel_2d_parallel is None: 

783 if cp is None: 

784 raise ImportError("CuPy is required for GPU Sobel operations") 

785 

786 _sobel_2d_parallel = cp.ElementwiseKernel( 

787 'raw T input, int32 Y, int32 X, int32 mode, T cval', 

788 'T output', 

789 ''' 

790 // Calculate which slice, row, col this thread handles 

791 int z = i / (Y * X); 

792 int y = (i % (Y * X)) / X; 

793 int x = i % X; 

794 

795 // Calculate base index for this slice 

796 int slice_base = z * Y * X; 

797 

798 // Helper function to get pixel with configurable boundary handling 

799 auto get_pixel = [&](int py, int px) -> T { 

800 if (mode == 0) { // constant 

801 if (py < 0 || py >= Y || px < 0 || px >= X) return cval; 

802 } else if (mode == 1) { // reflect 

803 py = py < 0 ? -py : (py >= Y ? 2*Y - py - 2 : py); 

804 px = px < 0 ? -px : (px >= X ? 2*X - px - 2 : px); 

805 } else if (mode == 2) { // nearest 

806 py = py < 0 ? 0 : (py >= Y ? Y-1 : py); 

807 px = px < 0 ? 0 : (px >= X ? X-1 : px); 

808 } else if (mode == 3) { // wrap 

809 py = py < 0 ? py + Y : (py >= Y ? py - Y : py); 

810 px = px < 0 ? px + X : (px >= X ? px - X : px); 

811 } 

812 return input[slice_base + py * X + px]; 

813 }; 

814 

815 // Sobel X kernel: [[-1,0,1],[-2,0,2],[-1,0,1]] (within slice only) 

816 T gx = -get_pixel(y-1, x-1) + get_pixel(y-1, x+1) + 

817 -2*get_pixel(y, x-1) + 2*get_pixel(y, x+1) + 

818 -get_pixel(y+1, x-1) + get_pixel(y+1, x+1); 

819 

820 // Sobel Y kernel: [[-1,-2,-1],[0,0,0],[1,2,1]] (within slice only) 

821 T gy = -get_pixel(y-1, x-1) - 2*get_pixel(y-1, x) - get_pixel(y-1, x+1) + 

822 get_pixel(y+1, x-1) + 2*get_pixel(y+1, x) + get_pixel(y+1, x+1); 

823 

824 // Calculate magnitude 

825 output = sqrt(gx*gx + gy*gy); 

826 ''', 

827 'sobel_2d_parallel' 

828 ) 

829 return _sobel_2d_parallel 

830 

831@cupy_func 

832def sobel_2d_vectorized(image: "cp.ndarray", mode: str = "reflect", cval: float = 0.0) -> "cp.ndarray": 

833 """ 

834 Apply 2D Sobel edge detection to all slices simultaneously - TRUE GPU PARALLELIZED. 

835 

836 Each slice is treated as an independent 2D grayscale image. All pixels across 

837 all slices are processed in parallel on the GPU with slice independence guaranteed. 

838 Uses ElementwiseKernel for maximum performance and true parallelization. 

839 

840 Args: 

841 image: 3D CuPy array of shape (Z, Y, X) 

842 mode: Boundary handling mode ('constant', 'reflect', 'nearest', 'wrap') 

843 cval: Constant value for 'constant' mode 

844 

845 Returns: 

846 Edge magnitude as 3D CuPy array of shape (Z, Y, X) 

847 """ 

848 _validate_3d_array(image) 

849 

850 # Map mode string to integer 

851 mode_map = {'constant': 0, 'reflect': 1, 'nearest': 2, 'wrap': 3} 

852 if mode not in mode_map: 

853 raise ValueError(f"Unknown mode: {mode}. Valid modes: {list(mode_map.keys())}") 

854 

855 mode_int = mode_map[mode] 

856 Z, Y, X = image.shape 

857 input_float = image.astype(cp.float32) 

858 output = cp.zeros_like(input_float) 

859 

860 # Launch parallel kernel - each thread processes one pixel 

861 sobel_kernel = _get_sobel_2d_kernel() 

862 sobel_kernel(input_float, Y, X, mode_int, cp.float32(cval), output) 

863 

864 return output.astype(image.dtype) 

865 

866@cupy_func 

867def sobel_3d_voxel(image: "cp.ndarray", mode: str = "reflect", cval: float = 0.0) -> "cp.ndarray": 

868 """ 

869 Apply true 3D voxel Sobel edge detection including Z-axis gradients - GPU PARALLELIZED. 

870 

871 This computes gradients in all three spatial dimensions (X, Y, Z) for true 

872 volumetric edge detection. Useful for 3D structure analysis where Z-dimension 

873 has spatial meaning (not just independent slices). 

874 

875 Args: 

876 image: 3D CuPy array of shape (Z, Y, X) 

877 mode: Boundary handling mode ('constant', 'reflect', 'nearest', 'wrap', 'mirror') 

878 cval: Constant value for 'constant' mode 

879 

880 Returns: 

881 3D edge magnitude as 3D CuPy array of shape (Z, Y, X) 

882 """ 

883 _validate_3d_array(image) 

884 

885 # Convert to float32 for processing 

886 image_float = image.astype(cp.float32) 

887 

888 # Apply Sobel filters in all three directions - GPU PARALLELIZED 

889 sobel_x = ndimage.sobel(image_float, axis=2, mode=mode, cval=cval) # X-direction gradients 

890 sobel_y = ndimage.sobel(image_float, axis=1, mode=mode, cval=cval) # Y-direction gradients 

891 sobel_z = ndimage.sobel(image_float, axis=0, mode=mode, cval=cval) # Z-direction gradients 

892 

893 # Calculate 3D magnitude - GPU PARALLELIZED 

894 magnitude = cp.sqrt(sobel_x**2 + sobel_y**2 + sobel_z**2) 

895 

896 return magnitude.astype(image.dtype) 

897 

898@cupy_func