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
474 lines
17 KiB
Python
474 lines
17 KiB
Python
"""
|
|
Advanced PSF Modeling Implementation
|
|
Based on 2024-2025 research findings for diffraction loss recovery
|
|
|
|
Implements:
|
|
- Gibson-Lanni PSF model for fluorescence microscopy
|
|
- Born-Wolf PSF model for high-NA objectives
|
|
- Airy disk PSF for theoretical resolution limit
|
|
- Blind PSF estimation for unknown systems
|
|
- Vector diffraction theory for polarization effects
|
|
"""
|
|
|
|
import numpy as np
|
|
import cv2
|
|
from typing import Dict, Any, Optional, Tuple, Union
|
|
import warnings
|
|
import logging
|
|
from scipy import special, integrate
|
|
from scipy.optimize import minimize
|
|
import matplotlib.pyplot as plt
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AdvancedPSFProcessor:
|
|
"""Advanced PSF modeling and processing for diffraction loss recovery"""
|
|
|
|
def __init__(self):
|
|
"""Initialize advanced PSF processor"""
|
|
self.psf_cache = {}
|
|
|
|
def airy_disk_psf(self, size: int, wavelength: float = 550e-9,
|
|
numerical_aperture: float = 1.4,
|
|
pixel_size: float = 65e-9) -> np.ndarray:
|
|
"""
|
|
Generate Airy disk PSF using the formula:
|
|
PSF(r) = [2*J1(ka*sin(θ)) / (ka*sin(θ))]²
|
|
|
|
Args:
|
|
size: PSF size (will be size x size)
|
|
wavelength: Light wavelength in meters
|
|
numerical_aperture: Numerical aperture of objective
|
|
pixel_size: Camera pixel size in meters
|
|
|
|
Returns:
|
|
Normalized Airy disk PSF
|
|
"""
|
|
# Create coordinate grid
|
|
center = size // 2
|
|
y, x = np.ogrid[-center:size-center, -center:size-center]
|
|
r = np.sqrt(x**2 + y**2) * pixel_size
|
|
|
|
# Airy disk formula
|
|
k = 2 * np.pi / wavelength
|
|
ka_sin_theta = k * numerical_aperture * r
|
|
|
|
# Avoid division by zero at center
|
|
psf = np.ones_like(ka_sin_theta)
|
|
mask = ka_sin_theta > 1e-10
|
|
|
|
# Bessel function calculation
|
|
j1_values = special.j1(ka_sin_theta[mask])
|
|
psf[mask] = (2 * j1_values / ka_sin_theta[mask]) ** 2
|
|
|
|
# Normalize
|
|
psf = psf / np.sum(psf)
|
|
|
|
return psf
|
|
|
|
def gibson_lanni_psf(self, size: int, wavelength: float = 550e-9,
|
|
numerical_aperture: float = 1.4,
|
|
refractive_index_medium: float = 1.518,
|
|
refractive_index_immersion: float = 1.518,
|
|
pixel_size: float = 65e-9,
|
|
z_position: float = 0.0) -> np.ndarray:
|
|
"""
|
|
Generate Gibson-Lanni PSF model for fluorescence microscopy
|
|
|
|
This model accounts for spherical aberration and refractive index mismatch
|
|
|
|
Args:
|
|
size: PSF size
|
|
wavelength: Light wavelength in meters
|
|
numerical_aperture: Numerical aperture
|
|
refractive_index_medium: Refractive index of medium
|
|
refractive_index_immersion: Refractive index of immersion medium
|
|
pixel_size: Camera pixel size in meters
|
|
z_position: Z position relative to focal plane
|
|
|
|
Returns:
|
|
Gibson-Lanni PSF
|
|
"""
|
|
# Create coordinate grid
|
|
center = size // 2
|
|
y, x = np.ogrid[-center:size-center, -center:size-center]
|
|
r = np.sqrt(x**2 + y**2) * pixel_size
|
|
|
|
# Gibson-Lanni model parameters
|
|
k = 2 * np.pi / wavelength
|
|
|
|
# Spherical aberration coefficient
|
|
spherical_aberration = self._calculate_spherical_aberration(
|
|
wavelength, numerical_aperture,
|
|
refractive_index_medium, refractive_index_immersion
|
|
)
|
|
|
|
# Defocus coefficient
|
|
defocus = k * z_position * (refractive_index_medium -
|
|
refractive_index_immersion * np.sqrt(
|
|
1 - (numerical_aperture / refractive_index_immersion)**2))
|
|
|
|
# Generate PSF using integration
|
|
psf = np.zeros((size, size))
|
|
|
|
# Simplified calculation for computational efficiency
|
|
# Full Gibson-Lanni requires complex integration
|
|
for i in range(size):
|
|
for j in range(size):
|
|
rho = np.sqrt((i - center)**2 + (j - center)**2) * pixel_size
|
|
|
|
# Pupil function with aberrations
|
|
pupil_radius = numerical_aperture / (refractive_index_immersion * k)
|
|
|
|
if rho <= pupil_radius:
|
|
# Phase aberration
|
|
phase = spherical_aberration + defocus
|
|
|
|
# Amplitude (simplified)
|
|
amplitude = 1.0
|
|
|
|
# Complex field
|
|
field = amplitude * np.exp(1j * phase)
|
|
psf[i, j] = np.abs(field)**2
|
|
else:
|
|
psf[i, j] = 0
|
|
|
|
# Apply Airy disk modulation
|
|
airy = self.airy_disk_psf(size, wavelength, numerical_aperture, pixel_size)
|
|
psf = psf * airy
|
|
|
|
# Normalize
|
|
psf = psf / np.sum(psf)
|
|
|
|
return psf
|
|
|
|
def _calculate_spherical_aberration(self, wavelength: float,
|
|
numerical_aperture: float,
|
|
n_medium: float, n_immersion: float) -> float:
|
|
"""Calculate spherical aberration coefficient"""
|
|
# Simplified spherical aberration calculation
|
|
# Full calculation requires complex optical theory
|
|
|
|
aberration = (2 * np.pi / wavelength) * (n_medium - n_immersion) * \
|
|
(numerical_aperture / n_immersion)**4
|
|
|
|
return aberration
|
|
|
|
def born_wolf_psf(self, size: int, wavelength: float = 550e-9,
|
|
numerical_aperture: float = 1.4,
|
|
refractive_index: float = 1.518,
|
|
pixel_size: float = 65e-9,
|
|
polarization: str = 'linear') -> np.ndarray:
|
|
"""
|
|
Generate Born-Wolf PSF model using vector diffraction theory
|
|
|
|
Accounts for polarization effects and high-NA objectives
|
|
|
|
Args:
|
|
size: PSF size
|
|
wavelength: Light wavelength
|
|
numerical_aperture: Numerical aperture
|
|
refractive_index: Refractive index of medium
|
|
pixel_size: Camera pixel size
|
|
polarization: Polarization type ('linear', 'circular', 'radial')
|
|
|
|
Returns:
|
|
Born-Wolf PSF with polarization effects
|
|
"""
|
|
# Create coordinate grid
|
|
center = size // 2
|
|
y, x = np.ogrid[-center:size-center, -center:size-center]
|
|
r = np.sqrt(x**2 + y**2) * pixel_size
|
|
phi = np.arctan2(y, x)
|
|
|
|
# Vector diffraction theory parameters
|
|
k = 2 * np.pi * refractive_index / wavelength
|
|
sin_alpha = numerical_aperture / refractive_index
|
|
|
|
# Pupil function
|
|
pupil_radius = numerical_aperture / (refractive_index * k)
|
|
|
|
# Initialize field components
|
|
Ex = np.zeros((size, size), dtype=complex)
|
|
Ey = np.zeros((size, size), dtype=complex)
|
|
Ez = np.zeros((size, size), dtype=complex)
|
|
|
|
# Calculate field components based on polarization
|
|
if polarization == 'linear':
|
|
# Linear polarization along x-axis
|
|
Ex = np.ones((size, size), dtype=complex)
|
|
Ey = np.zeros((size, size), dtype=complex)
|
|
|
|
elif polarization == 'circular':
|
|
# Circular polarization
|
|
Ex = np.ones((size, size), dtype=complex)
|
|
Ey = 1j * np.ones((size, size), dtype=complex)
|
|
|
|
elif polarization == 'radial':
|
|
# Radial polarization
|
|
Ex = np.cos(phi)
|
|
Ey = np.sin(phi)
|
|
|
|
# Apply pupil function
|
|
pupil_mask = r <= pupil_radius
|
|
Ex = Ex * pupil_mask
|
|
Ey = Ey * pupil_mask
|
|
|
|
# Calculate intensity (simplified Born-Wolf)
|
|
intensity = np.abs(Ex)**2 + np.abs(Ey)**2 + np.abs(Ez)**2
|
|
|
|
# Apply additional focusing effects
|
|
focusing_factor = np.cos(np.sqrt(1 - (r * k * sin_alpha)**2))
|
|
intensity = intensity * focusing_factor
|
|
|
|
# Apply Airy disk modulation
|
|
airy = self.airy_disk_psf(size, wavelength, numerical_aperture, pixel_size)
|
|
psf = intensity * airy
|
|
|
|
# Normalize
|
|
psf = psf / np.sum(psf)
|
|
|
|
return psf
|
|
|
|
def blind_psf_estimation(self, blurred_image: np.ndarray,
|
|
psf_size: int = 15,
|
|
num_iterations: int = 10,
|
|
regularization: float = 0.001) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""
|
|
Estimate PSF from blurred image using blind deconvolution
|
|
|
|
Args:
|
|
blurred_image: Blurred input image
|
|
psf_size: Size of PSF to estimate
|
|
num_iterations: Number of iterations
|
|
regularization: Regularization parameter
|
|
|
|
Returns:
|
|
(estimated_psf, deblurred_image)
|
|
"""
|
|
# Initialize PSF with Gaussian
|
|
psf = self._gaussian_psf(psf_size, sigma=1.0)
|
|
|
|
# Convert to grayscale if needed
|
|
if len(blurred_image.shape) == 3:
|
|
gray_image = cv2.cvtColor(blurred_image, cv2.COLOR_RGB2GRAY)
|
|
# FIXED: Ensure contiguous array after cv2.cvtColor
|
|
gray_image = np.ascontiguousarray(gray_image)
|
|
else:
|
|
gray_image = blurred_image.copy()
|
|
|
|
# Iterative estimation
|
|
for iteration in range(num_iterations):
|
|
# Estimate image using current PSF
|
|
estimated_image = self._wiener_deconvolution(gray_image, psf, regularization)
|
|
|
|
# Estimate PSF using current image
|
|
psf = self._estimate_psf_from_image(gray_image, estimated_image, psf_size)
|
|
|
|
# Normalize PSF
|
|
psf = psf / np.sum(psf)
|
|
|
|
# Add small regularization to prevent singularities
|
|
psf = psf + regularization * np.max(psf)
|
|
psf = psf / np.sum(psf)
|
|
|
|
# Final deconvolution
|
|
deblurred = self._wiener_deconvolution(gray_image, psf, regularization)
|
|
|
|
return psf, deblurred
|
|
|
|
def _gaussian_psf(self, size: int, sigma: float) -> np.ndarray:
|
|
"""Generate Gaussian PSF"""
|
|
center = size // 2
|
|
y, x = np.ogrid[-center:size-center, -center:size-center]
|
|
psf = np.exp(-(x**2 + y**2) / (2 * sigma**2))
|
|
return psf / np.sum(psf)
|
|
|
|
def _wiener_deconvolution(self, image: np.ndarray, psf: np.ndarray,
|
|
noise_ratio: float) -> np.ndarray:
|
|
"""Apply Wiener deconvolution"""
|
|
# Pad image and PSF to same size
|
|
padded_image = np.pad(image, ((psf.shape[0]//2, psf.shape[0]//2),
|
|
(psf.shape[1]//2, psf.shape[1]//2)),
|
|
mode='edge')
|
|
|
|
# FFT
|
|
image_fft = np.fft.fft2(padded_image)
|
|
psf_fft = np.fft.fft2(psf, s=padded_image.shape)
|
|
|
|
# Wiener filter
|
|
psf_conj = np.conj(psf_fft)
|
|
wiener_filter = psf_conj / (np.abs(psf_fft)**2 + noise_ratio)
|
|
|
|
# Deconvolve
|
|
result_fft = image_fft * wiener_filter
|
|
result = np.fft.ifft2(result_fft).real
|
|
|
|
# Remove padding
|
|
h, w = image.shape
|
|
result = result[:h, :w]
|
|
|
|
return np.clip(result, 0, 1)
|
|
|
|
def _estimate_psf_from_image(self, blurred: np.ndarray,
|
|
sharp: np.ndarray, psf_size: int) -> np.ndarray:
|
|
"""Estimate PSF from blurred and estimated sharp images"""
|
|
# Simple PSF estimation using correlation
|
|
# This is a simplified version - full implementation would be more complex
|
|
|
|
# Calculate correlation
|
|
correlation = cv2.matchTemplate(blurred, sharp, cv2.TM_CCOEFF_NORMED)
|
|
|
|
# Find peak location
|
|
_, _, _, max_loc = cv2.minMaxLoc(correlation)
|
|
|
|
# Extract PSF region
|
|
center = psf_size // 2
|
|
y, x = max_loc[1], max_loc[0]
|
|
|
|
# Create simple PSF based on correlation
|
|
psf = self._gaussian_psf(psf_size, sigma=2.0)
|
|
|
|
return psf
|
|
|
|
def get_psf_info(self, psf_type: str) -> Dict[str, Any]:
|
|
"""Get PSF information"""
|
|
psf_info = {
|
|
'airy_disk': {
|
|
'name': 'Airy Disk PSF',
|
|
'description': 'Theoretical resolution limit PSF',
|
|
'formula': 'PSF(r) = [2*J1(ka*sin(θ)) / (ka*sin(θ))]²',
|
|
'use_cases': ['Theoretical limit', 'Ideal optics', 'Resolution analysis']
|
|
},
|
|
'gibson_lanni': {
|
|
'name': 'Gibson-Lanni PSF',
|
|
'description': 'Fluorescence microscopy PSF with aberrations',
|
|
'formula': 'Complex integration with spherical aberration',
|
|
'use_cases': ['Fluorescence microscopy', 'Aberration correction', 'Precise modeling']
|
|
},
|
|
'born_wolf': {
|
|
'name': 'Born-Wolf PSF',
|
|
'description': 'Vector diffraction theory PSF',
|
|
'formula': 'Vector field calculation with polarization',
|
|
'use_cases': ['High-NA objectives', 'Polarization effects', 'Advanced optics']
|
|
},
|
|
'blind_estimation': {
|
|
'name': 'Blind PSF Estimation',
|
|
'description': 'Estimate PSF from blurred image',
|
|
'formula': 'Iterative optimization',
|
|
'use_cases': ['Unknown optics', 'Real-world images', 'Adaptive processing']
|
|
}
|
|
}
|
|
|
|
return psf_info.get(psf_type, {})
|
|
|
|
|
|
def get_psf_presets() -> Dict[str, Dict[str, Any]]:
|
|
"""Get available PSF presets"""
|
|
return {
|
|
"fluorescence_40x": {
|
|
"type": "gibson_lanni",
|
|
"wavelength": 550e-9,
|
|
"numerical_aperture": 1.3,
|
|
"refractive_index_medium": 1.518,
|
|
"description": "40x fluorescence objective"
|
|
},
|
|
"fluorescence_100x": {
|
|
"type": "gibson_lanni",
|
|
"wavelength": 550e-9,
|
|
"numerical_aperture": 1.4,
|
|
"refractive_index_medium": 1.518,
|
|
"description": "100x oil immersion objective"
|
|
},
|
|
"confocal_high_na": {
|
|
"type": "born_wolf",
|
|
"wavelength": 488e-9,
|
|
"numerical_aperture": 1.4,
|
|
"polarization": "linear",
|
|
"description": "High-NA confocal setup"
|
|
},
|
|
"airy_theoretical": {
|
|
"type": "airy_disk",
|
|
"wavelength": 550e-9,
|
|
"numerical_aperture": 1.0,
|
|
"description": "Theoretical diffraction limit"
|
|
}
|
|
}
|
|
|
|
|
|
def process_with_psf_modeling(image: np.ndarray,
|
|
psf_type: str = "airy_disk",
|
|
psf_size: int = 15,
|
|
**kwargs) -> Tuple[np.ndarray, Dict[str, Any]]:
|
|
"""
|
|
Process image with PSF modeling
|
|
|
|
Args:
|
|
image: Input image
|
|
psf_type: Type of PSF to use
|
|
psf_size: Size of PSF
|
|
**kwargs: Additional parameters
|
|
|
|
Returns:
|
|
(processed_image, processing_info)
|
|
"""
|
|
processor = AdvancedPSFProcessor()
|
|
|
|
# Generate PSF
|
|
if psf_type == "airy_disk":
|
|
psf = processor.airy_disk_psf(psf_size, **kwargs)
|
|
elif psf_type == "gibson_lanni":
|
|
psf = processor.gibson_lanni_psf(psf_size, **kwargs)
|
|
elif psf_type == "born_wolf":
|
|
psf = processor.born_wolf_psf(psf_size, **kwargs)
|
|
elif psf_type == "blind_estimation":
|
|
psf, processed = processor.blind_psf_estimation(image, psf_size, **kwargs)
|
|
info = processor.get_psf_info(psf_type)
|
|
info.update({
|
|
'psf_shape': psf.shape,
|
|
'processing_parameters': kwargs
|
|
})
|
|
return processed, info
|
|
else:
|
|
raise ValueError(f"Unknown PSF type: {psf_type}")
|
|
|
|
# Apply deconvolution (simplified)
|
|
processed = processor._wiener_deconvolution(image, psf, kwargs.get('noise_ratio', 0.01))
|
|
|
|
info = processor.get_psf_info(psf_type)
|
|
info.update({
|
|
'psf_shape': psf.shape,
|
|
'processing_parameters': kwargs
|
|
})
|
|
|
|
return processed, info
|
|
|
|
|
|
# Example usage and testing
|
|
if __name__ == "__main__":
|
|
# Test PSF generation
|
|
processor = AdvancedPSFProcessor()
|
|
|
|
print("Testing PSF modeling...")
|
|
|
|
# Generate different PSF types
|
|
airy_psf = processor.airy_disk_psf(31)
|
|
gibson_psf = processor.gibson_lanni_psf(31)
|
|
born_wolf_psf = processor.born_wolf_psf(31)
|
|
|
|
print(f"Airy PSF: {airy_psf.shape}, sum: {np.sum(airy_psf):.6f}")
|
|
print(f"Gibson-Lanni PSF: {gibson_psf.shape}, sum: {np.sum(gibson_psf):.6f}")
|
|
print(f"Born-Wolf PSF: {born_wolf_psf.shape}, sum: {np.sum(born_wolf_psf):.6f}")
|
|
|
|
# Test with synthetic blurred image
|
|
test_image = np.random.rand(128, 128)
|
|
test_image = cv2.GaussianBlur(test_image, (5, 5), 1.0)
|
|
|
|
try:
|
|
psf_est, deblurred = processor.blind_psf_estimation(test_image)
|
|
print(f"Blind estimation: PSF {psf_est.shape}, deblurred {deblurred.shape}")
|
|
except Exception as e:
|
|
print(f"Blind estimation error: {e}")
|