Pytorch图像去噪实战（五）：FFDNet可控图像去噪实战，用噪声强度图解决不同噪声等级问题

Pytorch图像去噪实战五FFDNet可控图像去噪实战用噪声强度图解决不同噪声等级问题一、问题场景同一个模型面对不同噪声强度时效果不稳定前面几篇我们做了 DnCNN、UNet、ResUNet 和 Attention UNet。这些模型在固定噪声强度下效果不错但在真实项目中我遇到了一个很麻烦的问题用户上传的图片噪声强度不一致同一个模型无法稳定处理所有情况。比如有些图只是轻微压缩噪声有些图是低光高ISO噪声有些图是截图二次压缩噪声有些图是扫描件颗粒噪声如果模型训练时只见过 sigma25那么面对 sigma10 时可能过度去噪面对 sigma50 时又去不干净。我一开始的做法是训练多个模型sigma15 一个模型sigma25 一个模型sigma50 一个模型但工程上非常麻烦模型多推理复杂部署成本高噪声估计不准时效果会崩因此这一篇我们做一个更工程化的方案FFDNet。二、FFDNet解决什么问题FFDNet的核心思想是不只输入带噪图像还额外输入一张噪声强度图。也就是说模型输入不再是noisy_image而是noisy_image noise_level_map这样模型就知道当前图片大概有多脏从而控制去噪力度。三、为什么噪声强度图很有用普通模型的问题是模型自己猜噪声强度。FFDNet的问题建模方式是我们直接告诉模型噪声强度。比如sigma15轻度去噪sigma25中度去噪sigma50强力去噪这在工程中非常实用因为我们可以根据业务场景动态调整去噪强度。四、工程目录结构ffdnet_denoise/ ├── data/ │ ├── train/ │ └── val/ ├── models/ │ └── ffdnet.py ├── dataset.py ├── train.py ├── eval.py └── utils.py五、数据集构建返回噪声强度图FFDNet训练时需要返回三个数据noisy带噪图sigma_map噪声强度图clean干净图dataset.pyimportosimportrandomimporttorchfromPILimportImagefromtorch.utils.dataimportDatasetimporttorchvision.transformsastransformsclassFFDNetDataset(Dataset):def__init__(self,root_dir,patch_size128):self.paths[os.path.join(root_dir,name)fornameinos.listdir(root_dir)ifname.lower().endswith((.jpg,.png,.jpeg))]self.patch_sizepatch_size self.to_tensortransforms.ToTensor()def__len__(self):returnlen(self.paths)def__getitem__(self,idx):imgImage.open(self.paths[idx]).convert(L)w,himg.sizeifwself.patch_sizeandhself.patch_size:xrandom.randint(0,w-self.patch_size)yrandom.randint(0,h-self.patch_size)imgimg.crop((x,y,xself.patch_size,yself.patch_size))else:imgimg.resize((self.patch_size,self.patch_size))cleanself.to_tensor(img)sigmarandom.uniform(0,50)sigma_valuesigma/255.0noisetorch.randn_like(clean)*sigma_value noisytorch.clamp(cleannoise,0.0,1.0)sigma_maptorch.ones_like(clean)*sigma_valuereturnnoisy,sigma_map,clean六、FFDNet模型实现这里我们实现一个简化版 FFDNet。核心是把 noisy 和 sigma_map 在通道维度拼接xtorch.cat([noisy,sigma_map],dim1)models/ffdnet.pyimporttorchimporttorch.nnasnnclassFFDNet(nn.Module):def__init__(self,in_channels2,out_channels1,features64):super().__init__()layers[]layers.append(nn.Conv2d(in_channels,features,3,padding1))layers.append(nn.ReLU(inplaceTrue))for_inrange(10):layers.append(nn.Conv2d(features,features,3,padding1))layers.append(nn.BatchNorm2d(features))layers.append(nn.ReLU(inplaceTrue))layers.append(nn.Conv2d(features,out_channels,3,padding1))self.netnn.Sequential(*layers)defforward(self,noisy,sigma_map):xtorch.cat([noisy,sigma_map],dim1)residualself.net(x)returnnoisy-residual七、训练代码train.pyimporttorchfromtorch.utils.dataimportDataLoaderfromdatasetimportFFDNetDatasetfrommodels.ffdnetimportFFDNetdeftrain():devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)datasetFFDNetDataset(data/train)loaderDataLoader(dataset,batch_size16,shuffleTrue,num_workers4)modelFFDNet().to(device)optimizertorch.optim.Adam(model.parameters(),lr1e-4)criteriontorch.nn.L1Loss()forepochinrange(1,51):model.train()total_loss0fornoisy,sigma_map,cleaninloader:noisynoisy.to(device)sigma_mapsigma_map.to(device)cleanclean.to(device)predmodel(noisy,sigma_map)losscriterion(pred,clean)optimizer.zero_grad()loss.backward()optimizer.step()total_lossloss.item()print(fEpoch{epoch}, Loss:{total_loss/len(loader):.6f})ifepoch%100:torch.save(model.state_dict(),fffdnet_epoch_{epoch}.pth)if__name____main__:train()八、推理时如何控制去噪强度这正是 FFDNet 最实用的地方。如果你觉得图片噪声轻就给小 sigmasigma15/255.0如果噪声很重就给大 sigmasigma50/255.0完整推理代码importtorchfromPILimportImageimporttorchvision.transformsastransformsimporttorchvision.utilsasvutilsfrommodels.ffdnetimportFFDNet devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)modelFFDNet().to(device)model.load_state_dict(torch.load(ffdnet_epoch_50.pth,map_locationdevice))model.eval()imgImage.open(test.png).convert(L)to_tensortransforms.ToTensor()noisyto_tensor(img).unsqueeze(0).to(device)sigma25/255.0sigma_maptorch.ones_like(noisy)*sigmawithtorch.no_grad():predmodel(noisy,sigma_map)predtorch.clamp(pred,0.0,1.0)vutils.save_image(pred.cpu(),denoised.png)九、为什么训练时sigma要随机如果只训练固定 sigma比如 25那么模型只会处理一种噪声。FFDNet的优势来自sigmarandom.uniform(0,50)这让模型在训练时见到连续噪声强度从而学会根据 sigma_map 控制去噪力度。十、真实工程中的使用方式在实际系统中可以提供三个模式轻度去噪sigma10适合轻微压缩噪声。标准去噪sigma25适合普通截图、扫描件。强力去噪sigma50适合低光、颗粒明显的图片。这样用户可以根据视觉效果选择强度工程体验会比单一模型好很多。十一、踩坑记录坑1sigma_map范围写错sigma_map必须和图像一样归一化到 0~1。错误写法sigma_maptorch.ones_like(clean)*25正确写法sigma_maptorch.ones_like(clean)*25/255.0坑2推理时忘记输入sigma_mapFFDNet不是普通单输入模型推理时必须传两个输入predmodel(noisy,sigma_map)坑3sigma给太大导致图像过度平滑如果 sigma80可能会把真实纹理也当噪声去掉。建议控制在0 ~ 50十二、效果验证FFDNet的优势不是单点PSNR最高而是同一个模型可以适配多个噪声强度。实际效果表现噪声强度普通UNetFFDNetsigma15容易过度去噪更自然sigma25效果接近稳定sigma50去噪不足更干净十三、适合收藏总结FFDNet完整流程输入带噪图构造噪声强度图noisy 与 sigma_map 拼接模型预测残差噪声noisy - residual 得到结果避坑清单sigma_map必须归一化训练时sigma要随机推理必须传sigma_mapsigma过大会过度平滑FFDNet适合做可控去噪十四、优化建议可以继续优化把基础网络换成UNet加残差块加注意力模块对真实噪声做噪声估计结合图像质量评分自动选择sigma结尾总结FFDNet真正有价值的地方在于工程可控性。很多时候图像去噪不是追求一个固定模型处理所有情况而是需要根据噪声强度灵活调整。FFDNet提供了一个很实用的思路把噪声强度显式告诉模型让模型按需去噪。下一篇预告Pytorch图像去噪实战六CBDNet真实噪声去噪实战解决合成噪声到真实噪声的泛化问题