"""Torch-accelerated Richardson-Lucy deconvolution.""" from __future__ import annotations import math from typing import Optional, Tuple import numpy as np try: # Optional dependency import torch import torch.nn.functional as F except ImportError as exc: # pragma: no cover - runtime guard torch = None # type: ignore _TORCH_IMPORT_ERROR = exc else: _TORCH_IMPORT_ERROR = None def _select_device(device_preference: str) -> "torch.device": if torch is None: # pragma: no cover - guard for static analyzers raise RuntimeError("PyTorch is required for GPU Richardson-Lucy") from _TORCH_IMPORT_ERROR pref = (device_preference or "auto").lower() if pref == "auto": if torch.cuda.is_available(): return torch.device("cuda") mps_backend = getattr(torch.backends, "mps", None) if mps_backend and mps_backend.is_available(): return torch.device("mps") return torch.device("cpu") try: device = torch.device(pref) except (TypeError, ValueError): print(f"Unknown device '{device_preference}', falling back to CPU") return torch.device("cpu") if device.type == "cuda" and not torch.cuda.is_available(): print("CUDA requested but not available, falling back to CPU") return torch.device("cpu") if device.type == "mps": mps_backend = getattr(torch.backends, "mps", None) if not (mps_backend and mps_backend.is_available()): print("MPS requested but not available, falling back to CPU") return torch.device("cpu") return device def _generate_psf_tensor( blur_type: str, blur_size: float, motion_angle: float, motion_length: float, device: "torch.device", dtype: "torch.dtype", ) -> "torch.Tensor": blur_type = blur_type.lower() if blur_type == "gaussian": sigma = max(blur_size, 1e-3) size = int(2 * math.ceil(3 * sigma) + 1) size = max(size, 3) coords = torch.arange(size, device=device, dtype=dtype) coords = coords - coords.mean() yy, xx = torch.meshgrid(coords, coords, indexing="ij") psf = torch.exp(-((xx ** 2 + yy ** 2) / (2 * sigma ** 2))) elif blur_type == "motion": length = max(1, int(round(motion_length))) size = max(length, 3) if size % 2 == 0: size += 1 psf = torch.zeros((size, size), device=device, dtype=dtype) center = size // 2 angle_rad = math.radians(motion_angle) dx = math.cos(angle_rad) * (length - 1) / 2.0 dy = math.sin(angle_rad) * (length - 1) / 2.0 x0, y0 = center - dx, center - dy x1, y1 = center + dx, center + dy for step in range(length): t = step / max(length - 1, 1) x = round(x0 + (x1 - x0) * t) y = round(y0 + (y1 - y0) * t) if 0 <= x < size and 0 <= y < size: psf[y, x] = 1.0 else: raise ValueError(f"Unsupported blur type '{blur_type}' for GPU Richardson-Lucy") psf_sum = psf.sum() if psf_sum <= 0: psf.fill_(0) psf[psf.shape[0] // 2, psf.shape[1] // 2] = 1.0 psf_sum = psf.sum() return psf / psf_sum def _depthwise_conv(image: "torch.Tensor", kernel: "torch.Tensor") -> "torch.Tensor": pad_h = kernel.shape[-2] // 2 pad_w = kernel.shape[-1] // 2 padded = F.pad(image, (pad_w, pad_w, pad_h, pad_h), mode="reflect") return F.conv2d(padded, kernel, groups=image.shape[1]) def _apply_laplacian_regularization( image: "torch.Tensor", strength: float, epsilon: float ) -> "torch.Tensor": if strength <= 0: return image laplacian_kernel = torch.tensor( [[0.0, 1.0, 0.0], [1.0, -4.0, 1.0], [0.0, 1.0, 0.0]], device=image.device, dtype=image.dtype, ).view(1, 1, 3, 3) laplacian_kernel = laplacian_kernel.repeat(image.shape[1], 1, 1, 1) laplacian = _depthwise_conv(image, laplacian_kernel) return torch.clamp(image - strength * laplacian, min=epsilon) def richardson_lucy_deconvolution_gpu( image: np.ndarray, psf: Optional[np.ndarray] = None, iterations: int = 10, clip: bool = True, filter_epsilon: Optional[float] = None, blur_type: str = "gaussian", blur_size: float = 2.0, motion_angle: float = 0.0, motion_length: float = 10.0, regularization: float = 0.0, device: str = "auto", precision: str = "fp32", ) -> np.ndarray: """Accelerated Richardson-Lucy deconvolution using PyTorch.""" if torch is None: raise RuntimeError("PyTorch is required for GPU Richardson-Lucy") from _TORCH_IMPORT_ERROR if image.ndim not in (2, 3): raise ValueError("Image must be 2D or 3D array") if iterations < 1: raise ValueError("Iterations must be >= 1") image = np.ascontiguousarray(image) alpha: Optional[np.ndarray] = None if image.ndim == 3 and image.shape[2] == 4: alpha = image[..., 3:] image = image[..., :3] input_float = image.astype(np.float32) / 255.0 if input_float.ndim == 2: input_tensor = torch.from_numpy(input_float).unsqueeze(0).unsqueeze(0) else: input_tensor = torch.from_numpy(input_float).permute(2, 0, 1).unsqueeze(0) target_device = _select_device(device) compute_dtype = torch.float16 if precision == "fp16" and target_device.type == "cuda" else torch.float32 if compute_dtype == torch.float16 and target_device.type != "cuda": compute_dtype = torch.float32 observed = input_tensor.to(device=target_device, dtype=compute_dtype) estimate = observed.clone() eps = filter_epsilon if filter_epsilon is not None else torch.finfo(observed.dtype).eps if psf is not None: if psf.ndim != 2: raise ValueError("Custom PSF must be 2D") psf_tensor = torch.from_numpy(np.ascontiguousarray(psf.astype(np.float32))) psf_tensor = psf_tensor.to(device=target_device, dtype=compute_dtype) psf_tensor = psf_tensor / psf_tensor.sum().clamp(min=eps) else: psf_tensor = _generate_psf_tensor( blur_type, blur_size, motion_angle, motion_length, target_device, compute_dtype ) kernel = psf_tensor.view(1, 1, psf_tensor.shape[0], psf_tensor.shape[1]) kernel = kernel.repeat(observed.shape[1], 1, 1, 1) kernel_flipped = torch.flip(kernel, dims=(-2, -1)) for _ in range(iterations): convolved = _depthwise_conv(estimate, kernel) convolved = torch.clamp(convolved, min=eps) ratio = observed / convolved correlation = _depthwise_conv(ratio, kernel_flipped) estimate = estimate * correlation estimate = _apply_laplacian_regularization(estimate, regularization, eps) estimate = torch.clamp(estimate, min=eps) if clip: estimate = torch.clamp(estimate, 0.0, 1.0) estimate = estimate.squeeze(0) if estimate.ndim == 3: result = estimate.permute(1, 2, 0).to(torch.float32).cpu().numpy() result = np.clip(result * 255.0 + 0.5, 0, 255).astype(np.uint8) if alpha is not None: result = np.concatenate([result, alpha], axis=2) else: result = estimate.to(torch.float32).cpu().numpy() result = np.clip(result * 255.0 + 0.5, 0, 255).astype(np.uint8) return np.ascontiguousarray(result) __all__ = ["richardson_lucy_deconvolution_gpu"]