ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

Unet天文图像降噪实战:从噪声模型到PyTorch实现

2026/9/11 19:40:42 拓冰建站 浏览量
Unet天文图像降噪实战:从噪声模型到PyTorch实现 简介基于Unet网络实现天文图像降噪的源码与项目说明包面向计算机、人工智能等专业学生及深度学习入门者适合作为课程设计、毕业设计或图像去噪方向的项目起步参考。项目采用先为图像叠加噪声、再以Unet网络将带噪图像映射回原始图像的方式进行训练数据生成、模型搭建与训练流程均完整提供。资源内含209个文件以90个npy数据文件、91个png图像样本为主配合3个Python脚本、2个ipynb教程/结果文件以及README说明文档整体压缩包约35.94MB目录结构清晰便于按数据处理、模型训练和结果展示模块查阅。已有383人浏览学习。代码均在测试运行成功后上传可在Google Colab平台直接运行使用generate_data.py生成训练数据并通过ipynb完成Unet模型的训练与效果对比随附项目说明有助快速理解各文件作用便于二次改进和迁移到其他图像降噪场景。1. 天文图像的噪声模型为什么降噪偏偏绕不开 Unet一张叠加了 30 次曝光的仙女座星系照片单帧读噪声可能在 5e- 左右叠加后随机噪声确实会降低但暗弱的旋臂结构往往依然淹没在背景方差里。天文降噪和普通手机拍照降噪最大的差异是噪声不是简单的高斯白噪声而是泊松噪声与读出噪声的混合体而且信号本身的高动态范围让“一刀切”的滤波核完全失效。传统的均值滤波、中值滤波、小波阈值去噪在平滑噪声的同时必然抹掉恒星点扩散函数PSF的细节甚至把双星洗成单星。深度学习里的 Unet 因为其编码器-解码器结构和跳跃连接天生适合做“保持形状的同时去掉噪声”这种任务这就是本项目的核心思路。如果你手头有一批 FITS 格式的天文照片希望在不依赖专业软件的情况下自己训练一个针对性降噪模型或者你只是想知道 Unet 作为图像分割经典网络怎么迁移到降噪任务上这篇文章就是给你写的。整个方案只需要 Python 和 PyTorch数据清洗、噪声模拟、模型训练、验证流程都会逐步展开。下面我会先讲清楚结构选型的原因再给出一套可以直接跑的代码框架最后聊几个真正影响降噪效果、但论文里很少写的细节。2. Unet 做天文降噪的结构逻辑编码器、跳跃连接与损失函数设计2.1 天文降噪任务里 Unet 的编码器-解码器到底在学什么Unet 的编码器是逐步下采样的卷积栈每下采样一次特征图分辨率减半、通道数翻倍。这个设计对天文图像非常友好因为星云和星系的自相似结构可以用多尺度特征描述。编码器浅层学到的是高频边缘和点状源的锐利程度深层学到的是大尺度背景梯度。但问题在于如果只用编码器提取特征再直接上采样很多暗弱细节会丢失这正是当年全连接分割网络的问题。解码器的工作是逐级恢复空间分辨率它需要把深层的高层语义和浅层的边缘信息融合。值得注意的是天文降噪和语义分割不同——分割任务是输出类别掩码而降噪任务输出的是连续像素值。所以解码器最后一层的激活函数一般是线性或者不经过激活直接输出估计的干净像素。在我的实现里输入和输出尺寸保持一致所以上采样路径不需要额外裁剪这也是 Unet 做图像恢复比做分割更顺手的原因之一。这里有一个关键点Unet 并不是唯一的降噪架构DnCNN、FFDNet 也常见。但 Unet 的优势在于它天然支持输入输出 shape 不变的端到端训练而且跳跃连接让梯度可以更顺畅地流向浅层。对于天文图像这种信噪比空间分布不均匀的数据浅层特征的保留尤为重要。我一般不会一开始就换 Attention U-Net 或 Res-UNet除非基线效果不理想。2.2 跳跃连接是保细节的关键参数上怎么配合跳跃连接把编码器第 i 层的输出直接拼接到解码器对应层。对于降噪任务这一步的意义是提供原始分辨率的参考信息相当于给解码器一条“旁路”让它不需要从零重建细节。我在实践中发现如果把跳跃连接去掉模型虽然也能收敛但星点轮廓会变模糊PSF 的半峰全宽FWHM可能从 2.5 像素被抹成 4 像素。这就是为什么很多改进版本的 Unet 会引入注意力机制去控制跳跃连接的权重而不是无脑拼接。参数上跳跃连接涉及一个容易忽略的点输入图像的归一化范围。天文图像往往是 16 位 FITS数值范围从 0 到 65535如果直接输入网络卷积权重初始化就崩了。我习惯先把数据线性拉伸到 [0,1] 范围。但有一个细节如果只按全图最大最小拉伸那么极亮的恒星会压制暗弱的细节导致噪声模型不完整。更好的做法是先裁剪一定百分位的上下限比如用 2% 和 98% 分位数做线性拉伸然后再统一减均值、除方差。这样跳跃连接拼接的两层特征分布不会出现量级失衡。2.3 损失函数选 L2 还是 SSIM结合泊松噪声的加权方案大部分降噪网络默认用 L2 损失也就是均方误差。L2 损失在优化时梯度平滑PSNR 指标也直接对应均方误差所以用 L2 训练时 PSNR 通常涨得更快。但 L2 的缺点是容易产生过度平滑的结果因为它在像素级对预测取平均会抹掉细节中的高频波动。另一种做法是组合损失L_total L_mse λ * (1 - SSIM_loss)。SSIM 从亮度、对比度、结构三方面衡量感知质量它对局部结构的保留比 L2 强。不过天文图像上直接使用 SSIM 有一个坑——SSIM 对局部窗口内的均值和方差敏感而天文图像里星点所在区域的局部方差很大SSIM 梯度会把这些区域当成重要区域来优化反而可能让背景噪点被保留。我自己的方案是采用带噪声模型的损失加权因为泊松噪声的信号方差等于信号强度在暗弱区域方差小在亮区方差大。因此可以使用变分近似对每个像素的 MSE 按置信度加权。下面是一个简化的带权损失代码import torch import torch.nn.functional as F def poisson_weighted_mse(pred, target, exposure1.0): # 目标图像近似作为泊松噪声的期望值 variance torch.clamp(target * exposure 1e-6, min1e-6) weights 1.0 / torch.sqrt(variance) diff (pred - target) ** 2 return (weights * diff).mean() def combined_loss(pred, target, lambda_ssim0.1): mse F.mse_loss(pred, target) ssim_val ssim_loss(pred, target) # 这里用你选择的SSIM实现返回 1-ssim return mse lambda_ssim * ssim_val逻辑说明poisson_weighted_mse中我们把target当作真实信号的估计用target的强度估算噪声方差。曝光时间exposure用于控制泊松噪声的相对大小。理论上亮区权重小暗区权重大这样模型不会只盯着亮核而是把暗弱细节也拉起来。实际训练时我通常选用加权 MSE 加一个很小的 SSIM 项而不是单用 L2。参数建议exposure根据你的数据决定。如果数据是单帧短曝光取值 1.0 左右如果是叠加后的图像噪声方差已经降低可把exposure设为 0.2~0.5否则暗区权重会过大。lambda_ssim建议从 0.05 开始调SSIM 项太大会导致梯度方向偏向结构相似训练初期容易震荡。3. 用 Python 搭一个可运行的天文图像降噪 Unet 最小项目3.1 环境准备与依赖这个项目只需要基础的深度学习库不需要安装天文专业耗时依赖。Python 版本建议 3.8 或 3.10PyTorch 2.x 都可以。如果你还没有配置环境直接创建虚拟环境python -m venv astro_denoise_env source astro_denoise_env/bin/activate # Windows 下用 astro_denoise_env\Scripts\activate pip install torch torchvision numpy astropy tensorboard tqdm这里把astropy加进来了它用来读取 FITS 文件。有些天文图像是 .fits 格式里面包含 WCS 坐标信息和噪声参数用astropy.io.fits读取最标准。注意torchvision可以帮助我们处理图像变换但天文 FITS 不一定能被ImageFolder直接读取所以数据加载一般用自定义 Dataset。环境里最容易出问题的是 PyTorch 版本和 CUDA 版本不匹配。我用纯 CPU 训练小尺寸图像也可以只是慢。建议先用 16x16 或 32x32 的裁剪块调试代码确认无误后再上 GPU。3.2 数据加载与噪声注入泊松噪声模拟代码真实天文图像配对数据很难获取常见做法是用高信噪比图像加模拟噪声来构造训练对。高信噪比图可以来自哈勃数据、SDSS 数据或者你自己的长曝光叠加图。以下代码展示如何把干净图像变成带泊松噪声的退化图import numpy as np from astropy.io import fits def add_poisson_gaussian_noise(clean_img, gain1.0, read_noise10.0): clean_img: 归一化到 [0, 1] 的 float32 图像 gain: 电子增益 (e-/ADU) read_noise: 读噪声标准差 (e-) # 将归一化的强度转换回光子数这里假设 clean_img 已经乘了一个参考满阱值 # 更合理的做法是直接用原始 e- 计数这里为示例做线性映射 electrons clean_img * 65535.0 / gain noisy_electrons np.random.poisson(electrons).astype(np.float32) noisy_electrons np.random.normal(0, read_noise, sizeelectrons.shape).astype(np.float32) # 转换回 [0,1] 范围并做截断 noisy_img np.clip(noisy_electrons * gain / 65535.0, 0.0, 1.0) return noisy_img逻辑说明先假设输入的clean_img代表归一化后的真实信号乘以一个虚拟满阱电荷数这里用 65535映射到电子数目。泊松分布生成的随机数就是带信号依赖噪声的电子数然后加一个高斯分布的读噪声最后再除回去。这样生成的噪声模型与真实 CCD 接近。实际上真正的暗场噪声是非稳态的但作为训练数据这种模拟已经足够让网络学会区分散粒噪声和高斯噪声。对训练数据我建议不要每次都重新计算噪声而是把干净图像保存成.npy数组在训练循环里在线生成噪声。这样每个 epoch 噪声都会重新采样等效于无限多的训练样本对提升泛化性非常有帮助。3.3 Unet 模型定义面向图像恢复的紧凑版经典的 Unet 最开始用于 256x256 输入但天文图像可能非常大为了避免显存爆炸输入通常裁剪成 64x64 或 128x128。下面定义了一个适合降噪任务的轻型 Unet重点是把通道数安排得比分割任务小一些因为降噪任务不需要太深的语义信息import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNetDenoise(nn.Module): def __init__(self, in_channels1, out_channels1, features[32, 64, 128, 256]): super().__init__() self.downs nn.ModuleList() self.ups nn.ModuleList() self.pool nn.MaxPool2d(2) # 编码器 for f in features: self.downs.append(DoubleConv(in_channels, f)) in_channels f self.bottleneck DoubleConv(features[-1], features[-1]*2) # 解码器 for f in reversed(features): self.ups.append(nn.ConvTranspose2d(f*2, f, kernel_size2, stride2)) self.ups.append(DoubleConv(f*2, f)) self.final_conv nn.Conv2d(features[0], out_channels, kernel_size1) def forward(self, x): skip_connections [] for down in self.downs: x down(x) skip_connections.append(x) x self.pool(x) x self.bottleneck(x) skip_connections skip_connections[::-1] for idx in range(0, len(self.ups), 2): x self.ups[idx](x) skip skip_connections[idx//2] if x.shape ! skip.shape: x nn.functional.interpolate(x, sizeskip.shape[2:], modebilinear, align_cornersTrue) x torch.cat((skip, x), dim1) x self.ups[idx1](x) return self.final_conv(x)代码说明DoubleConv是标准的两层 3x3 卷积使用 BatchNorm 和 ReLU。features列表控制了各层通道数从 32 逐步翻倍到 256。这个模型参数量约 8M对小尺寸天文图块足够。解码器中的ConvTranspose2d是转置卷积用来上采样。skip_connections保存编码器各层特征并在解码器中用torch.cat拼接。注意这个实现里有一个潜在的分辨率错位问题当输入尺寸不是 2 的幂时池化后的尺寸向下取整上采样后的尺寸可能和跳跃连接不一致。我在代码里加入了interpolate来对齐但这样会破坏一些空间对应关系。最好在数据加载时把图像统一裁剪成 64x64 或 128x128从根本上避免该问题。3.4 训练循环与超参数batch size 怎么影响收敛下面是一个最小训练循环包含模型初始化、优化器和日志输出。训练时我通常使用 AdamW 而不是 Adam因为配合权重衰减可以让网络更稳定。device torch.device(cuda if torch.cuda.is_available() else cpu) model UNetDenoise(in_channels1, out_channels1).to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100, eta_min1e-6) def train_one_epoch(loader, model, optimizer, criterion): model.train() total_loss 0 for clean_batch in loader: clean_batch clean_batch.to(device) # 在线加噪声 noisy_batch add_noise_batch(clean_batch) # 你需要把噪声函数向量化或使用 numpy 处理 noisy_batch torch.from_numpy(noisy_batch).float().to(device) pred model(noisy_batch) loss criterion(pred, clean_batch) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader)超参数选择逻辑batch size 小的时候 BatchNorm 的统计量不稳定容易导致训练震荡。显存允许的情况下尽量用 32 以上。如果只能跑 8那就把 BatchNorm 换成 GroupNorm或者直接在DoubleConv里去掉 BatchNorm。我测试过在天文图像上 GroupNorm 数量从 8 到 16 都有效但 BatchNorm 在 batch size 小于 8 时确实会掉点。学习率方面1e-4 是起步值配合 Cosine 衰减在 100 个 epoch 内逐渐降到 1e-6。如果发现 loss 不下降先把学习率调小一个量级而不是盲目增加模型容量。3.5 训练中监控与断点保存训练时不仅要看 loss还要在每个 epoch 结束把模型权重保存下来。我习惯保存两份一份是最新 epoch 的last.pth一份是验证集 PSNR 最高的best.pth。代码如下best_psnr 0 for epoch in range(epochs): train_loss train_one_epoch(...) val_psnr validate(model, val_loader) if val_psnr best_psnr: best_psnr val_psnr torch.save(model.state_dict(), best_unet_denoise.pth) if epoch % 10 0: torch.save(model.state_dict(), fcheckpoint_epoch{epoch}.pth) scheduler.step()validate函数需要在无梯度模式下计算 PSNR并且注意预测输出可能要截断到 0-1 再算指标。保存模型后下次训练可以直接加载权重继续跑。4. 训练调参实战如何把 PSNR 和 SSIM 指标稳定提上去4.1 预处理比网络结构更重要减暗场、平场和归一化实际天文图像拿到手之后不能直接丢进网络。CCD 图像有偏置电平bias、暗电流dark和像素响应不均匀flat。这些系统性偏差如果不校准Unet 学到的“噪声”里就会混入固定模式噪声导致在真实数据集上泛化变差。标准的校准流程是:from astropy.io import fits import numpy as np def calibrate(fits_path, bias_path, dark_path, flat_path): with fits.open(fits_path) as img_hdu, fits.open(bias_path) as bias_hdu, \ fits.open(dark_path) as dark_hdu, fits.open(flat_path) as flat_hdu: image img_hdu[0].data.astype(np.float32) bias bias_hdu[0].data.astype(np.float32) dark dark_hdu[0].data.astype(np.float32) flat flat_hdu[0].data.astype(np.float32) calibrated (image - bias - dark) / np.maximum(flat, 1.0) # 去除异常像素做插值 bad_mask ~np.isfinite(calibrated) calibrated[bad_mask] np.nanmedian(calibrated) return calibrated这段代码里flat要防止除零所以用np.maximum把最小值限制到 1。暗场可能已经包含偏置需要根据设备类型判断。做完校准后再执行裁剪和归一化。如果跳过这步网络会花大量容量去学偏置模式的周期条纹而不是真正的噪声。4.2 数据增强策略裁剪、翻转与旋转天文图像和自然图像不同物体方向没有“上下”概念所以旋转和翻转增强可以放心用。但注意翻滚和旋转后图像中的星点 PSF 各向同性所以不会引入伪影。我常用的增强如下import random import numpy as np def augment(image): # 随机90度旋转 if random.random() 0.5: image np.rot90(image, krandom.randint(1, 3)) if random.random() 0.5: image np.flip(image, axis0) if random.random() 0.5: image np.flip(image, axis1) return np.ascontiguousarray(image)注意增强是在干净图和噪声图上共用一套随机变换否则配对被破坏。另外不要用随机亮度抖动和色彩抖动因为天文图像的亮度是有物理含义的亮度扰动会破坏泊松噪声的比例关系。4.3 训练阶段的三个关键参数学习率调度、梯度裁剪和指数移动平均除了学习率梯度裁剪对 Unet 训练稳定性很有帮助。天文图像中亮恒星区域的梯度可能非常大如果不裁剪很容易导致 loss 炸成 nan。PyTorch 里一行代码torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)另一个实用技巧是使用指数移动平均EMA对模型参数做滑动平均。EMA 可以在训练末期提升 PSNR 0.2~0.5 dB。实现方式不复杂维护一个影子权重每次更新后按照衰减系数更新影子权重最后用影子权重做验证。学习率调度建议用 Cosine 或 ReduceLROnPlateau。如果使用 Cosine初始学习率可以稍微大一点因为后面会自动降下去。如果使用 ReduceLROnPlateau它的参数patience建议设为 8~10避免验证集指标微小波动导致误降。4.4 容易过拟合的迹象与应对策略天文图像数据量通常不大一个深度 Unet 很容易在训练集上 PSNR 很高验证集上效果普通。过拟合的典型表现是训练 loss 持续下降验证 loss 在第 30 个 epoch 左右开始反弹。此时最优先的做法是降低模型复杂度把features从[32,64,128,256]降为[16,32,64,128]参数量会缩小 4 倍但降噪能力可能只下降 0.3dB。这比硬加 Dropout 更可靠。另一种有效做法是 MixUp 或 CutMix 样式的数据混合但这个在图像恢复任务里不太常用。我更推荐使用预训练权重迁移学习。如果实在没有预训练权重就先在小尺寸裁剪块上训练再利用上采样初始化大尺寸模型这个过程称为渐进式训练。另外需要注意的是验证集应该选用和训练集不同的天体区域而不是随机打乱。否则网络可能记住了固定星点的位置导致验证指标虚高。一般按星表坐标切分让训练集和验证集没有重叠的显著恒星。5. 验证技巧用噪声功率谱和残差图判断降噪质量而不是只看 PSNR5.1 PSNR 和 SSIM 的盲区很多人在验证的时候只看 PSNR但 PSNR 是全局像素误差的统计量它无法告诉我们噪声是否在空间上被均匀地抑制。我在实践中遇到过这样的情况模型把亮星周围的噪声消掉了但背景区域出现规则的棋盘格伪影PSNR 依然提升了 2dB。SSIM 对局部结构敏感但对低频噪声反应比较迟缓。所以最终验证阶段一定要回到天文图像本身的特点用噪声功率谱来检查是否存在“过度平滑”或“伪纹理”。噪声功率谱的计算方式取一块没有明显星源的背景区域对残差图输出图像减原始干净图像做二维 FFT再对功率谱做径向平均。如果降噪后残差的功率谱在高频段仍接近原始噪声说明模型没有有效抑制高频噪声如果功率谱在某个频率出现尖峰说明模型产生了周期伪影。import numpy as np def radial_power_spectrum(img): # 输入二维图像返回径向平均功率谱 fshift np.fft.fftshift(np.fft.fft2(img - img.mean())) power np.abs(fshift) ** 2 y, x np.indices(power.shape) cx, cy power.shape[1] // 2, power.shape[0] // 2 r np.sqrt((x - cx) ** 2 (y - cy) ** 2).astype(np.int) tbin np.bincount(r.ravel(), weightspower.ravel()) nr np.bincount(r.ravel()) radial tbin / np.maximum(nr, 1) return radial这个函数返回一个一维数组代表以中心频率为原点的径向平均功率。使用的时候对原噪声图残差和降噪后残差分别画曲线理想的结果是降噪后的曲线整体向下移动而不是出现局部凸起。5.2 残差图与径向亮度剖面残差图的意义在于直观显示系统残留结构。把干净图和输出图相减检查残差图中是否还有恒星的形状。如果残差在恒星位置呈现环形结构说明 Unet 对 PSF 中心强度的估计有偏差。这时候可以考虑在损失函数中加入结构项但更好的办法是检查预处理阶段是否做过星点对齐。如果输入的图像存在亚像素平移Unet 很难同时保持位置精度和边缘锐度。验证效果时我最常做一个操作画一条穿过亮星中心和暗弱背景的直线对比原始噪声图、真实干净图和网络输出图沿这条线的亮度曲线。对于亮星看 FWHM 是否保持对于背景看波动范围是否明显缩小。这个步骤比任何指标都直观而且写报告时也更容易让人信服。5.3 把模型打包进项目说明中的注意事项如果你的最终交付物是python源代码项目说明.zip那么模型文件一般不要保存完整torch.save(model, ...)而是保存 state_dict。同时推荐导出为 TorchScript 或 ONNX方便后续在 C 或推理库中使用。导出 TorchScript 的常见坑是 BatchNorm 层在某些操作符上不支持动态尺寸所以固定输入尺寸最好。代码示例model.eval() example_input torch.rand(1, 1, 128, 128) traced_model torch.jit.trace(model, example_input) traced_model.save(unet_denoise_script.pt)注意trace方法会冻结控制流如果模型中包含interpolate的动态尺寸分支建议在导出前把输入固定为训练时尺寸。项目说明里要写清楚依赖版本、训练数据格式和最低显存要求避免别人拿到 zip 后卡在环境配置上。我一般会在说明文件里放一个requirements.txt并且写出 “CPU 环境下跑通完整训练流程至少需要 8GB 内存” 类似的话这样使用者心里有底。在最后还有一个容易被忽视的验证方法把输出图像保存成 FITS 文件在 DS9 里通过拉伸查看暗弱天区的噪声纹理。如果背景区域看起来像纯粹的高斯噪声说明降噪成功如果看到 coherent 的纹理或假结构就要检查数据增强和归一化步骤。验证阶段多花 10 分钟看细节胜过在指标排行榜上多追 0.1dB。本文还有配套的精品资源点击获取