PyTorch实战：用傅里叶变换给你的图片做‘体检’，分离振幅与相位（附完整代码）

PyTorch实战：用傅里叶变换给你的图片做‘体检’，分离振幅与相位（附完整代码）

当你拿到一张X光片时，医生能从中看出骨骼的健康状况。而傅里叶变换就是图像的"X光机"，它能将图像从空间域转换到频域，让我们看到隐藏在像素背后的频率特征。今天，我们就用PyTorch这把"手术刀"，来解剖图像的频率成分。

传统图像处理中，OpenCV和NumPy是傅里叶变换的常用工具。但在深度学习时代，PyTorch提供了更高效的torch.fft模块，能无缝集成到神经网络流程中。我们将通过一个生动的方式：分别固定振幅和相位，观察重建图像的变化，直观理解这两个关键成分对图像的贡献。

1. 准备工作：搭建图像"体检"环境

在开始前，我们需要准备好"体检"所需的工具。与医院体检需要各种仪器一样，我们的图像处理也需要一系列Python库的支持。

首先安装必要的依赖：

pip install torch torchvision matplotlib pillow

接下来导入我们将用到的模块：

import torch import torchvision.transforms as transforms from PIL import Image import matplotlib.pyplot as plt import numpy as np

提示：如果你使用GPU加速，确保安装对应CUDA版本的PyTorch。傅里叶变换这类计算密集型操作在GPU上会有显著性能提升。

2. 图像加载与预处理：准备"体检样本"

任何体检都需要采集样本，我们的样本就是待处理的图像。PyTorch处理图像有其特定的张量格式要求，我们需要将普通图像转换为适合处理的格式。

完整的图像加载和预处理代码如下：

def load_image(image_path, size=(512, 512)): # 使用PIL加载图像 image = Image.open(image_path).convert('RGB') # 定义转换流程 transform = transforms.Compose([ transforms.Resize(size), # 调整大小 transforms.ToTensor(), # 转换为张量 [0,1]范围 ]) # 应用转换 tensor = transform(image) # 添加批次维度 (C,H,W) -> (1,C,H,W) tensor = tensor.unsqueeze(0) return tensor # 使用示例 image_tensor = load_image("your_image.jpg") print(f"图像张量形状: {image_tensor.shape}")

这个预处理流程做了几件重要的事情：

1. 确保图像是RGB三通道格式
2. 将图像调整为统一尺寸（默认为512x512）
3. 转换为PyTorch张量并归一化到[0,1]范围
4. 添加批次维度以适应后续处理

3. 执行傅里叶变换：开始"影像检查"

现在，我们的图像已经躺在"检查床"上，准备接受傅里叶变换的"扫描"。PyTorch的torch.fft模块提供了高效的傅里叶变换实现。

完整的变换代码如下：

def fft_transform(image_tensor): # 对图像的最后两个维度(高和宽)执行傅里叶变换 fft = torch.fft.fftn(image_tensor, dim=(-2, -1)) # 将零频率分量移到频谱中心 fft_shifted = torch.fft.fftshift(fft, dim=(-2, -1)) return fft_shifted # 应用变换 fft_result = fft_transform(image_tensor)

理解这段代码的关键点：

• fftn是多维傅里叶变换，我们指定在最后两个维度（图像的高和宽）上执行变换
• fftshift将零频率分量移到频谱中心，这是可视化频谱时的常规操作
• 变换结果是复数张量，包含实部和虚部信息

4. 分离振幅与相位：分析"体检报告"

傅里叶变换的结果可以分解为两个关键部分：

• 振幅谱：表示各频率成分的强度
• 相位谱：表示各频率成分的位置信息

我们可以用以下代码分离这两个成分：

# 计算振幅谱和相位谱 amplitude = torch.abs(fft_result) phase = torch.angle(fft_result) # 可视化振幅谱（取对数便于观察） def visualize_spectrum(amplitude): # 取对数并归一化 log_amp = torch.log(amplitude + 1e-9) # 加小值避免log(0) log_amp = (log_amp - log_amp.min()) / (log_amp.max() - log_amp.min()) # 转换为numpy并显示 plt.imshow(log_amp.squeeze().permute(1, 2, 0).cpu().numpy()) plt.title('Amplitude Spectrum') plt.axis('off') plt.show() visualize_spectrum(amplitude)

5. 实验：振幅与相位谁更重要？

为了直观理解振幅和相位的作用，我们设计一个有趣的实验：分别固定其中一个变量，观察重建图像的变化。

5.1 固定相位，随机化振幅

def reconstruct_with_constant_phase(fft_result): # 获取原始相位 phase = torch.angle(fft_result) # 创建随机振幅（保持原始统计特性） original_amp = torch.abs(fft_result) random_amp = torch.rand_like(original_amp) * original_amp.mean() # 重建复数频谱 reconstructed = random_amp * torch.exp(1j * phase) # 逆傅里叶变换 image_recon = torch.fft.ifftn( torch.fft.ifftshift(reconstructed, dim=(-2, -1)), dim=(-2, -1) ) return torch.abs(image_recon) # 重建并显示图像 recon_image = reconstruct_with_constant_phase(fft_result) plt.imshow(recon_image.squeeze().permute(1, 2, 0).cpu().numpy()) plt.title('Reconstruction with Constant Phase') plt.axis('off') plt.show()

5.2 固定振幅，随机化相位

def reconstruct_with_constant_amplitude(fft_result): # 获取原始振幅 amplitude = torch.abs(fft_result) # 创建随机相位 random_phase = torch.rand_like(amplitude) * 2 * np.pi - np.pi # 重建复数频谱 reconstructed = amplitude * torch.exp(1j * random_phase) # 逆傅里叶变换 image_recon = torch.fft.ifftn( torch.fft.ifftshift(reconstructed, dim=(-2, -1)), dim=(-2, -1) ) return torch.abs(image_recon) # 重建并显示图像 recon_image = reconstruct_with_constant_amplitude(fft_result) plt.imshow(recon_image.squeeze().permute(1, 2, 0).cpu().numpy()) plt.title('Reconstruction with Constant Amplitude') plt.axis('off') plt.show()

实验结果对比：

这个实验清晰地表明：相位信息对图像内容的理解更为关键。即使振幅完全随机化，只要保留相位，图像的主要结构仍然可辨；而丢失相位信息后，图像内容完全无法识别。

6. 完整代码实现

以下是整合所有步骤的完整代码，你可以直接复制使用：

import torch import torchvision.transforms as transforms from PIL import Image import matplotlib.pyplot as plt import numpy as np def load_image(image_path, size=(512, 512)): image = Image.open(image_path).convert('RGB') transform = transforms.Compose([ transforms.Resize(size), transforms.ToTensor(), ]) tensor = transform(image).unsqueeze(0) return tensor def fft_transform(image_tensor): fft = torch.fft.fftn(image_tensor, dim=(-2, -1)) return torch.fft.fftshift(fft, dim=(-2, -1)) def reconstruct_image(fft_result): # 逆变换 image_recon = torch.fft.ifftn( torch.fft.ifftshift(fft_result, dim=(-2, -1)), dim=(-2, -1) ) return torch.abs(image_recon) def visualize_comparison(original, recon_phase, recon_amp): plt.figure(figsize=(15, 5)) # 原始图像 plt.subplot(1, 3, 1) plt.imshow(original.squeeze().permute(1, 2, 0).cpu().numpy()) plt.title('Original Image') plt.axis('off') # 固定相位重建 plt.subplot(1, 3, 2) plt.imshow(recon_phase.squeeze().permute(1, 2, 0).cpu().numpy()) plt.title('Constant Phase') plt.axis('off') # 固定振幅重建 plt.subplot(1, 3, 3) plt.imshow(recon_amp.squeeze().permute(1, 2, 0).cpu().numpy()) plt.title('Constant Amplitude') plt.axis('off') plt.tight_layout() plt.show() # 主流程 if __name__ == "__main__": # 1. 加载图像 image_tensor = load_image("your_image.jpg") # 2. 傅里叶变换 fft_result = fft_transform(image_tensor) # 3. 分离振幅和相位 amplitude = torch.abs(fft_result) phase = torch.angle(fft_result) # 4. 固定相位，随机振幅重建 random_amp = torch.rand_like(amplitude) * amplitude.mean() recon_phase = random_amp * torch.exp(1j * phase) image_phase = reconstruct_image(recon_phase) # 5. 固定振幅，随机相位重建 random_phase = torch.rand_like(phase) * 2 * np.pi - np.pi recon_amp = amplitude * torch.exp(1j * random_phase) image_amp = reconstruct_image(recon_amp) # 6. 可视化比较 visualize_comparison(image_tensor, image_phase, image_amp)

在实际项目中，我发现将傅里叶变换集成到PyTorch流程中有几个优势：GPU加速大幅提升计算效率；自动微分支持可以构建频域处理的神经网络；与torchvision的无缝集成简化了预处理流程。