如何正确替换FFT实部 复现CVPR2021傅里叶域数据增强方法

• 核心逻辑错误：误解了FFT复数的构成，论文要求的是固定幅值、保留原始相位，而你直接修改FFT的实部为常数，既没有实现固定幅值，还破坏了原始相位信息，这是结果差异大的最主要原因。复数的幅值计算公式为幅值=√(实部²+虚部²)，相位计算公式为相位=arctan2(虚部, 实部)，仅修改实部会同时改变幅值和相位，完全不符合论文的设计思路。
• 常数取值量级不匹配：你的输入图像已经被归一化到0_{1区间，此时FFT结果的幅值量级远小于你设置的200000，过大的常数会导致逆变换后的数值完全溢出合理范围，生成异常结果。论文里的20000是针对0}255范围的原始图像计算的幅值量级，你需要根据自己输入的归一化程度调整常数值，或者直接参考原始图像FFT幅值的平均量级设置。
• 维度变换错误：PyTorch处理的图像张量格式为[batch, channel, height, width]，你在逆变换后仅用transpose(2, 0)调整维度，会导致通道维度和空间维度顺序混乱，产生明显伪影。正确的维度调整应该用permute(1,2,0)将[channel, height, width]转换为OpenCV要求的[height, width, channel]格式。
• 颜色通道顺序错误：cv2.imread默认读取的是BGR格式的图像，而transforms.ToTensor()仅做数值范围和维度顺序转换，不会修正通道顺序，你最后用cv2.imshow输出RGB格式的张量时，会出现颜色错乱甚至伪影。

import cv2
import numpy as np
import torch
from torchvision import transforms

loader = transforms.Compose([transforms.ToTensor()])
def process_img(img_path, img_size):
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 提前转成RGB格式
    img = cv2.resize(img, (img_size, img_size))
    img = img.astype(np.float32) / 255.0
    img = loader(img)
    img = img.unsqueeze(0)
    return img

img = process_img("./data/house.jpg", 128)
img_fft = torch.fft.fft2(img, dim=(-2, -1))
# 正确提取幅值和相位
origin_amp = torch.abs(img_fft)
origin_phase = torch.angle(img_fft)
# 固定幅值，取值参考原始幅值的平均量级即可，不用硬套论文的20000
fixed_amp = torch.full_like(origin_amp, fill_value=origin_amp.mean())
# 用固定幅值+原始相位重建FFT结果
new_fft = fixed_amp * torch.exp(1j * origin_phase)
# 逆变换取实部
img_ifft = torch.fft.ifft2(new_fft, dim=(-2, -1)).real
# 调整维度到HWC格式
img_ifft = img_ifft.squeeze(0).permute(1,2,0).numpy()
# 数值裁剪到合理范围，转回BGR给OpenCV显示
img_ifft = np.clip(img_ifft, 0, 1)
img_ifft = cv2.cvtColor((img_ifft*255).astype(np.uint8), cv2.COLOR_RGB2BGR)
cv2.imshow("", img_ifft)
cv2.waitKey(0)

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