基于FFT的图像条纹噪声去除代码优化求助

核心问题拆解 #

图像褪色、去噪效果差的常见诱因：

• 误修改/滤除了频谱的DC分量（图像平均亮度），导致整体亮度失真
• 硬阈值掩码过于粗糙，要么没精准覆盖条纹噪声的频谱峰值，要么误删了有效信号的频谱
• 逆FFT后未正确将像素值映射回原范围，引发褪色
• 频谱突变导致振铃效应，让图像出现伪影

优化后的实现代码 #

import numpy as np
from scipy.signal import convolve

def remove_stripe_noise(img, stripe_direction='horizontal', sigma=2, smooth_kernel_size=3):
    # 兼容RGB/灰度图，RGB自动转灰度（如需保留色彩可对单通道分别处理）
    if len(img.shape) == 3:
        img = np.mean(img, axis=2).astype(np.float32)
    else:
        img = img.astype(np.float32)
    
    # 傅里叶变换+移频（将DC分量移至频谱中心）
    fft_img = np.fft.fft2(img)
    fft_shift = np.fft.fftshift(fft_img)
    
    # 提取频谱幅度，用于定位噪声峰值
    magnitude = np.abs(fft_shift)
    rows, cols = magnitude.shape
    center_row, center_col = rows // 2, cols // 2
    
    # 生成针对条纹方向的软掩码
    mask = np.ones_like(magnitude)
    if stripe_direction == 'horizontal':
        # 水平条纹对应频谱中的垂直亮线，统计列方向频谱分布
        col_profile = np.mean(magnitude, axis=0)
        # 筛选Top1%的峰值位置，排除DC点附近区域
        peak_cols = np.where(col_profile > np.percentile(col_profile, 99))[0]
        peak_cols = peak_cols[np.abs(peak_cols - center_col) > 5]
        
        # 用高斯衰减生成软掩码，避免硬截断的振铃效应
        for col in peak_cols:
            x = np.arange(cols)
            gauss_weight = np.exp(-((x - col)**2)/(2*sigma**2))
            mask[:, x] = np.minimum(mask[:, x], 1 - gauss_weight)
    else:
        # 垂直条纹对应频谱中的水平亮线，统计行方向频谱分布
        row_profile = np.mean(magnitude, axis=1)
        peak_rows = np.where(row_profile > np.percentile(row_profile, 99))[0]
        peak_rows = peak_rows[np.abs(peak_rows - center_row) > 5]
        
        for row in peak_rows:
            y = np.arange(rows)
            gauss_weight = np.exp(-((y - row)**2)/(2*sigma**2))
            mask[y, :] = np.minimum(mask[y, :], 1 - gauss_weight)
    
    # 强制保留DC分量，避免整体亮度失真
    mask[center_row, center_col] = 1.0
    
    # 应用掩码并逆傅里叶变换
    filtered_fft_shift = fft_shift * mask
    filtered_fft = np.fft.ifftshift(filtered_fft_shift)
    filtered_img = np.fft.ifft2(filtered_fft).real
    
    # 恢复像素值范围，解决褪色问题
    filtered_img = np.clip(filtered_img, np.min(img), np.max(img))
    filtered_img = (filtered_img - filtered_img.min()) / (filtered_img.max() - filtered_img.min()) * 255
    filtered_img = filtered_img.astype(np.uint8)
    
    # 空域平滑消除振铃伪影
    if smooth_kernel_size > 0:
        kernel = np.ones((smooth_kernel_size, smooth_kernel_size)) / (smooth_kernel_size**2)
        filtered_img = convolve(filtered_img, kernel, mode='same', boundary='symm')
        filtered_img = filtered_img.astype(np.uint8)
    
    return filtered_img

关键优化点说明 #

1. DC分量强制保留：直接锁定频谱中心的DC点不被掩码修改，从根源避免图像整体亮度下降（褪色）
2. 精准噪声定位：通过统计频谱的行列分布，筛选出条纹噪声对应的峰值区域，避免无差别滤波
3. 软掩码替代硬阈值：用高斯衰减曲线生成掩码，减少频谱突变带来的振铃伪影，让图像更自然
4. 像素范围严格恢复：逆FFT后将浮点值重新映射到0-255区间，确保输出图像的亮度和对比度与原图匹配
5. 空域后处理：用均值卷积核做轻微平滑，进一步消除FFT处理残留的细微振铃效应

使用示例 #

# 读取测试图像（替换为你的图像读取逻辑，比如用PIL或opencv）
from PIL import Image
test_img = np.array(Image.open("test_stripe.png"))
# 处理水平条纹噪声，可根据实际情况调整sigma和平滑核大小
denoised_img = remove_stripe_noise(test_img, stripe_direction='horizontal', sigma=3, smooth_kernel_size=3)
# 保存结果
Image.fromarray(denoised_img).save("denoised_result.png")

内容的提问来源于stack exchange，提问作者Maor Israeli