Eric Rollei
d6dd6470cf
Features: - 40+ image processing algorithms - Traditional denoising: Wavelet, NLM, BM3D (GPU-accelerated) - AI models: DnCNN, SwinIR, SCUNet, NAFNet, Noise-DA - DeepInv diffusion denoising integration - Advanced film grain detection and processing - Frequency domain processing - Richardson-Lucy deconvolution - Wiener filter - External DeepInv service with FastAPI - Comprehensive documentation and guides
307 lines
9.1 KiB
Python
307 lines
9.1 KiB
Python
"""
|
|
GPU-Accelerated BM3D Denoising using pytorch-bm3d
|
|
|
|
This module provides GPU-accelerated BM3D denoising using the pytorch-bm3d library.
|
|
Requires CUDA-capable GPU and compiled CUDA extension.
|
|
|
|
Performance: ~15-30x faster than CPU BM3D
|
|
- 256x256: ~0.08s (vs ~2.5s CPU)
|
|
- 1080p: ~0.14s (vs ~5-10s CPU)
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
# Try to import pytorch-bm3d
|
|
try:
|
|
from pytorch_bm3d import BM3D
|
|
PYTORCH_BM3D_AVAILABLE = True
|
|
except ImportError:
|
|
PYTORCH_BM3D_AVAILABLE = False
|
|
BM3D = None
|
|
|
|
# Try to import CUDA extension
|
|
try:
|
|
import bm3d_cuda
|
|
BM3D_CUDA_AVAILABLE = True
|
|
except ImportError:
|
|
BM3D_CUDA_AVAILABLE = False
|
|
|
|
|
|
def is_available():
|
|
"""
|
|
Check if GPU BM3D is available.
|
|
|
|
Returns:
|
|
tuple: (available, reason) where available is bool and reason is str
|
|
"""
|
|
if not PYTORCH_BM3D_AVAILABLE:
|
|
return False, "pytorch-bm3d not installed"
|
|
|
|
if not BM3D_CUDA_AVAILABLE:
|
|
return False, "CUDA extension not compiled"
|
|
|
|
if not torch.cuda.is_available():
|
|
return False, "CUDA not available"
|
|
|
|
return True, "GPU BM3D available"
|
|
|
|
|
|
def get_preset_parameters():
|
|
"""
|
|
Get preset parameter combinations for different noise types.
|
|
|
|
Returns:
|
|
dict: Preset name to parameters mapping
|
|
"""
|
|
return {
|
|
"light_noise": {
|
|
"sigma": 10,
|
|
"two_step": True,
|
|
"description": "Light noise (sigma=10), two-step refinement"
|
|
},
|
|
"medium_noise": {
|
|
"sigma": 25,
|
|
"two_step": True,
|
|
"description": "Medium noise (sigma=25), two-step refinement"
|
|
},
|
|
"heavy_noise": {
|
|
"sigma": 40,
|
|
"two_step": True,
|
|
"description": "Heavy noise (sigma=40), two-step refinement"
|
|
},
|
|
"fast_light": {
|
|
"sigma": 10,
|
|
"two_step": False,
|
|
"description": "Fast mode, light noise (single-step)"
|
|
},
|
|
"fast_medium": {
|
|
"sigma": 25,
|
|
"two_step": False,
|
|
"description": "Fast mode, medium noise (single-step)"
|
|
}
|
|
}
|
|
|
|
|
|
def bm3d_gpu_denoise(image, sigma=25, two_step=True, device=None):
|
|
"""
|
|
GPU-accelerated BM3D denoising.
|
|
|
|
Args:
|
|
image: Input image as numpy array [H, W, C] or [H, W], float [0-1] or uint8 [0-255]
|
|
sigma: Noise standard deviation (0-100). Represents noise level in 0-255 scale.
|
|
sigma=25 means "25/255 = 9.8% noise"
|
|
two_step: Use two-step refinement (higher quality, ~1.5x slower)
|
|
True = equivalent to 'refilter' profile
|
|
False = single-step (faster)
|
|
device: CUDA device (e.g., 'cuda:0' or torch.device). If None, uses cuda:0
|
|
|
|
Returns:
|
|
tuple: (denoised_image, info_dict)
|
|
- denoised_image: numpy array same shape as input, float [0-1]
|
|
- info_dict: Processing information
|
|
|
|
Raises:
|
|
RuntimeError: If GPU BM3D is not available
|
|
ValueError: If input parameters are invalid
|
|
"""
|
|
# Check availability
|
|
available, reason = is_available()
|
|
if not available:
|
|
raise RuntimeError(f"GPU BM3D not available: {reason}")
|
|
|
|
# Validate input
|
|
if not isinstance(image, np.ndarray):
|
|
raise ValueError("Image must be a numpy array")
|
|
|
|
if image.ndim not in [2, 3]:
|
|
raise ValueError(f"Image must be 2D or 3D, got shape {image.shape}")
|
|
|
|
if sigma < 0 or sigma > 100:
|
|
raise ValueError(f"Sigma must be in [0, 100], got {sigma}")
|
|
|
|
# Store original properties
|
|
original_shape = image.shape
|
|
original_dtype = image.dtype
|
|
is_grayscale = (image.ndim == 2)
|
|
|
|
# Convert to float [0-1] if needed
|
|
if image.dtype == np.uint8:
|
|
image_float = image.astype(np.float32) / 255.0
|
|
else:
|
|
image_float = image.astype(np.float32)
|
|
# Ensure [0-1] range
|
|
if image_float.max() > 1.0 or image_float.min() < 0.0:
|
|
image_float = np.clip(image_float, 0.0, 1.0)
|
|
|
|
# Handle grayscale
|
|
if is_grayscale:
|
|
image_float = image_float[:, :, np.newaxis] # [H, W] -> [H, W, 1]
|
|
|
|
h, w, c = image_float.shape
|
|
|
|
# Prepare info dictionary
|
|
info = {
|
|
"method": "BM3D-GPU",
|
|
"sigma": sigma,
|
|
"two_step": two_step,
|
|
"input_shape": original_shape,
|
|
"channels": c,
|
|
"device": str(device) if device else "cuda:0"
|
|
}
|
|
|
|
try:
|
|
# Set device
|
|
if device is None:
|
|
device = torch.device("cuda:0")
|
|
elif isinstance(device, str):
|
|
device = torch.device(device)
|
|
|
|
# Convert to pytorch-bm3d format: [1, C, H, W], int32 [0-255]
|
|
# NumPy: [H, W, C] float [0-1] -> PyTorch: [1, C, H, W] int [0-255]
|
|
image_uint8 = (image_float * 255.0).clip(0, 255).astype(np.uint8)
|
|
|
|
# Convert to torch tensor: [H, W, C] -> [C, H, W] -> [1, C, H, W]
|
|
image_torch = torch.from_numpy(image_uint8).permute(2, 0, 1).unsqueeze(0)
|
|
image_torch = image_torch.to(torch.int32).to(device)
|
|
|
|
# CRITICAL: Ensure contiguous memory layout
|
|
image_torch = image_torch.contiguous()
|
|
|
|
info["torch_shape"] = list(image_torch.shape)
|
|
info["torch_dtype"] = str(image_torch.dtype)
|
|
|
|
# Calculate variance from sigma
|
|
# variance = (sigma_in_255_scale)^2
|
|
variance = float(sigma ** 2)
|
|
info["variance"] = variance
|
|
|
|
# Initialize BM3D model
|
|
bm3d_model = BM3D(two_step=two_step)
|
|
|
|
# Run denoising with timing
|
|
torch.cuda.synchronize()
|
|
import time
|
|
start_time = time.time()
|
|
|
|
with torch.no_grad():
|
|
output_torch = bm3d_model(image_torch, variance=variance)
|
|
|
|
torch.cuda.synchronize()
|
|
elapsed = time.time() - start_time
|
|
|
|
info["processing_time"] = f"{elapsed:.3f}s"
|
|
|
|
# Convert back to numpy: [1, C, H, W] int [0-255] -> [H, W, C] float [0-1]
|
|
output_np = output_torch.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
|
output_np = output_np.astype(np.float32) / 255.0
|
|
output_np = np.clip(output_np, 0.0, 1.0)
|
|
|
|
# Restore original shape if grayscale
|
|
if is_grayscale:
|
|
output_np = output_np[:, :, 0] # [H, W, 1] -> [H, W]
|
|
|
|
# Calculate quality metrics if input wasn't too noisy
|
|
if sigma <= 50: # Only for reasonable noise levels
|
|
input_flat = image_float.flatten()
|
|
output_flat = output_np.flatten() if output_np.ndim == 3 else output_np[:, :, np.newaxis].flatten()
|
|
|
|
# Mean Squared Error
|
|
mse = np.mean((input_flat - output_flat[:len(input_flat)]) ** 2)
|
|
if mse > 0:
|
|
psnr = 10 * np.log10(1.0 / mse)
|
|
info["psnr"] = f"{psnr:.2f} dB"
|
|
|
|
# Clean up GPU memory
|
|
del image_torch, output_torch, bm3d_model
|
|
torch.cuda.empty_cache()
|
|
|
|
info["status"] = "success"
|
|
|
|
return output_np, info
|
|
|
|
except Exception as e:
|
|
info["status"] = "failed"
|
|
info["error"] = str(e)
|
|
|
|
# Clean up on error
|
|
torch.cuda.empty_cache()
|
|
|
|
raise RuntimeError(f"GPU BM3D denoising failed: {e}") from e
|
|
|
|
|
|
def get_optimal_sigma(noise_level_description):
|
|
"""
|
|
Get optimal sigma value based on noise description.
|
|
|
|
Args:
|
|
noise_level_description: String describing noise level
|
|
|
|
Returns:
|
|
int: Recommended sigma value
|
|
"""
|
|
noise_map = {
|
|
"very_light": 5,
|
|
"light": 10,
|
|
"medium": 25,
|
|
"heavy": 40,
|
|
"very_heavy": 60
|
|
}
|
|
|
|
return noise_map.get(noise_level_description.lower(), 25)
|
|
|
|
|
|
def estimate_performance(width, height):
|
|
"""
|
|
Estimate processing time for given image dimensions.
|
|
|
|
Args:
|
|
width: Image width
|
|
height: Image height
|
|
|
|
Returns:
|
|
dict: Performance estimates
|
|
"""
|
|
# Based on benchmark: 1080p (1920x1080 = 2,073,600 pixels) takes ~0.14s
|
|
pixels = width * height
|
|
base_pixels = 1920 * 1080
|
|
base_time = 0.14
|
|
|
|
# Roughly linear scaling with pixel count
|
|
estimated_time = (pixels / base_pixels) * base_time
|
|
|
|
# CPU equivalent (rough estimate: 20-30x slower)
|
|
estimated_cpu_time = estimated_time * 25
|
|
|
|
return {
|
|
"estimated_gpu_time": f"{estimated_time:.3f}s",
|
|
"estimated_cpu_time": f"{estimated_cpu_time:.1f}s",
|
|
"speedup": "~20-30x vs CPU",
|
|
"pixels": f"{pixels:,}"
|
|
}
|
|
|
|
|
|
# Convenience function for common use case
|
|
def denoise_image(image, noise_level="medium", quality="high"):
|
|
"""
|
|
Simplified interface for common denoising scenarios.
|
|
|
|
Args:
|
|
image: Input image numpy array
|
|
noise_level: "light", "medium", or "heavy"
|
|
quality: "high" (two-step) or "fast" (single-step)
|
|
|
|
Returns:
|
|
tuple: (denoised_image, info_dict)
|
|
"""
|
|
sigma_map = {
|
|
"light": 10,
|
|
"medium": 25,
|
|
"heavy": 40
|
|
}
|
|
|
|
sigma = sigma_map.get(noise_level.lower(), 25)
|
|
two_step = (quality.lower() == "high")
|
|
|
|
return bm3d_gpu_denoise(image, sigma=sigma, two_step=two_step)
|