AIRNet模型使用与代码分析(All-In-One Image Restoration Network)

2024-06-10 03:44

本文主要是介绍AIRNet模型使用与代码分析(All-In-One Image Restoration Network),希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

AIRNet提出了一种较为简易的pipeline,以单一网络结构应对多种任务需求(不同类型,不同程度)。但在效果上看,ALL-In-One是不如One-By-One的,且本文方法的亮点是batch内选择patch进行对比学习。在与sota对比上,仅是Denoise任务精度占优,在Derain与Dehaze任务上,效果不如One-By-One的MPRNet方法。本博客对AIRNet的关键结构实现,loss实现,data_patch实现进行深入分析,并对模型进行推理使用。

其论文的详细可以阅读:https://blog.csdn.net/a486259/article/details/139559389?spm=1001.2014.3001.5501

项目地址:https://blog.csdn.net/a486259/article/details/139559389?spm=1001.2014.3001.5501

项目依赖:torch、mmcv-full
安装mmcv-full时,需要注意torch所对应的cuda版本,要与系统中的cuda版本一致。

1、模型结构

AirNet的网络结构如下所示,输入图像x交由CBDE提取到嵌入空间z,z与x输入到DGRN模块的DGG block中逐步优化,最终输出预测结果。
在这里插入图片描述
模型代码在net\model.py

from torch import nnfrom net.encoder import CBDE
from net.DGRN import DGRNclass AirNet(nn.Module):def __init__(self, opt):super(AirNet, self).__init__()# Encoderself.E = CBDE(opt)  #编码特征值# Restorerself.R = DGRN(opt) #特征解码def forward(self, x_query, x_key):if self.training:fea, logits, labels, inter = self.E(x_query, x_key)restored = self.R(x_query, inter)return restored, logits, labelselse:fea, inter = self.E(x_query, x_query)restored = self.R(x_query, inter)return restored

1.1 CBDE模块

CBDE模块的功能是在模块内进行对比学习,核心是MoCo. Moco论文地址:https://arxiv.org/pdf/1911.05722

class CBDE(nn.Module):def __init__(self, opt):super(CBDE, self).__init__()dim = 256# Encoderself.E = MoCo(base_encoder=ResEncoder, dim=dim, K=opt.batch_size * dim)def forward(self, x_query, x_key):if self.training:# degradation-aware represenetion learningfea, logits, labels, inter = self.E(x_query, x_key)return fea, logits, labels, interelse:# degradation-aware represenetion learningfea, inter = self.E(x_query, x_query)return fea, inter

ResEncoder所对应的网络结构如下所示
在这里插入图片描述

在AIRNet中的CBDE模块里的MoCo模块的关键代码如下,其在内部自行完成了正负样本的分配,最终输出logits, labels用于计算对比损失的loss。但其所优化的模块实际上是ResEncoder。MoCo模块只是在训练阶段起作用,在推理阶段是不起作用的。

class MoCo(nn.Module):def forward(self, im_q, im_k):"""Input:im_q: a batch of query imagesim_k: a batch of key imagesOutput:logits, targets"""if self.training:# compute query featuresembedding, q, inter = self.encoder_q(im_q)  # queries: NxCq = nn.functional.normalize(q, dim=1)# compute key featureswith torch.no_grad():  # no gradient to keysself._momentum_update_key_encoder()  # update the key encoder_, k, _ = self.encoder_k(im_k)  # keys: NxCk = nn.functional.normalize(k, dim=1)# compute logits# Einstein sum is more intuitive# positive logits: Nx1l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)# negative logits: NxKl_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])# logits: Nx(1+K)logits = torch.cat([l_pos, l_neg], dim=1)# apply temperaturelogits /= self.T# labels: positive key indicatorslabels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()# dequeue and enqueueself._dequeue_and_enqueue(k)return embedding, logits, labels, interelse:embedding, _, inter = self.encoder_q(im_q)return embedding, inter

1.2 DGRN模块

DGRN模块的实现代码如下所示,可以看到核心是DGG模块,其不断迭代优化输入图像。

class DGRN(nn.Module):def __init__(self, opt, conv=default_conv):super(DGRN, self).__init__()self.n_groups = 5n_blocks = 5n_feats = 64kernel_size = 3# head modulemodules_head = [conv(3, n_feats, kernel_size)]self.head = nn.Sequential(*modules_head)# bodymodules_body = [DGG(default_conv, n_feats, kernel_size, n_blocks) \for _ in range(self.n_groups)]modules_body.append(conv(n_feats, n_feats, kernel_size))self.body = nn.Sequential(*modules_body)# tailmodules_tail = [conv(n_feats, 3, kernel_size)]self.tail = nn.Sequential(*modules_tail)def forward(self, x, inter):# headx = self.head(x)# bodyres = xfor i in range(self.n_groups):res = self.body[i](res, inter)res = self.body[-1](res)res = res + x# tailx = self.tail(res)return x

在这里插入图片描述
DGG模块的结构示意如下所示
在这里插入图片描述
DGG代码实现如下所示,DGG模块内嵌DGB模块,DGB模块内嵌DGM模块,DGM模块内嵌SFT_layer模块与DCN_layer(可变性卷积)
在这里插入图片描述

2、loss实现

AIRNet中提到的loss如下所示,其中Lrec是L1 loss,Lcl是Moco模块实现的对比损失。
在这里插入图片描述
AIRNet的loss实现代码在train.py中,CE loss是针对CBDE(Moco模块)的输出进行计算,l1 loss是针对修复图像与清晰图片。

    # Network Constructionnet = AirNet(opt).cuda()net.train()# Optimizer and Lossoptimizer = optim.Adam(net.parameters(), lr=opt.lr)CE = nn.CrossEntropyLoss().cuda()l1 = nn.L1Loss().cuda()# Start trainingprint('Start training...')for epoch in range(opt.epochs):for ([clean_name, de_id], degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2) in tqdm(trainloader):degrad_patch_1, degrad_patch_2 = degrad_patch_1.cuda(), degrad_patch_2.cuda()clean_patch_1, clean_patch_2 = clean_patch_1.cuda(), clean_patch_2.cuda()optimizer.zero_grad()if epoch < opt.epochs_encoder:_, output, target, _ = net.E(x_query=degrad_patch_1, x_key=degrad_patch_2)contrast_loss = CE(output, target)loss = contrast_losselse:restored, output, target = net(x_query=degrad_patch_1, x_key=degrad_patch_2)contrast_loss = CE(output, target)l1_loss = l1(restored, clean_patch_1)loss = l1_loss + 0.1 * contrast_loss# backwardloss.backward()optimizer.step()

这里可以看出来,AIRNet首先是训练CBDE模块,最后才训练CBDE模块+DGRN模块。

3、TrainDataset

TrainDataset的实现代码在utils\dataset_utils.py中,首先找到__getitem__函数进行分析。以下代码为关键部分,删除了大部分在逻辑上重复的部分。TrainDataset一共支持5种数据类型,‘denoise_15’: 0, ‘denoise_25’: 1, ‘denoise_50’: 2,是不需要图像对的(在代码里面自动对图像添加噪声);‘derain’: 3, ‘dehaze’: 4是需要图像对进行训练的。

class TrainDataset(Dataset):def __init__(self, args):super(TrainDataset, self).__init__()self.args = argsself.rs_ids = []self.hazy_ids = []self.D = Degradation(args)self.de_temp = 0self.de_type = self.args.de_typeself.de_dict = {'denoise_15': 0, 'denoise_25': 1, 'denoise_50': 2, 'derain': 3, 'dehaze': 4}self._init_ids()self.crop_transform = Compose([ToPILImage(),RandomCrop(args.patch_size),])self.toTensor = ToTensor()def __getitem__(self, _):de_id = self.de_dict[self.de_type[self.de_temp]]if de_id < 3:if de_id == 0:clean_id = self.s15_ids[self.s15_counter]self.s15_counter = (self.s15_counter + 1) % self.num_cleanif self.s15_counter == 0:random.shuffle(self.s15_ids)# clean_id = random.randint(0, len(self.clean_ids) - 1)clean_img = crop_img(np.array(Image.open(clean_id).convert('RGB')), base=16)clean_patch_1, clean_patch_2 = self.crop_transform(clean_img), self.crop_transform(clean_img)clean_patch_1, clean_patch_2 = np.array(clean_patch_1), np.array(clean_patch_2)# clean_name = self.clean_ids[clean_id].split("/")[-1].split('.')[0]clean_name = clean_id.split("/")[-1].split('.')[0]clean_patch_1, clean_patch_2 = random_augmentation(clean_patch_1, clean_patch_2)degrad_patch_1, degrad_patch_2 = self.D.degrade(clean_patch_1, clean_patch_2, de_id)clean_patch_1, clean_patch_2 = self.toTensor(clean_patch_1), self.toTensor(clean_patch_2)degrad_patch_1, degrad_patch_2 = self.toTensor(degrad_patch_1), self.toTensor(degrad_patch_2)self.de_temp = (self.de_temp + 1) % len(self.de_type)if self.de_temp == 0:random.shuffle(self.de_type)return [clean_name, de_id], degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2

可以看出TrainDataset返回的数据有:degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2。

3.1 clean_patch分析

通过以下代码可以看出 clean_patch_1, clean_patch_2是来自于同一个图片,然后基于crop_transform变化,变成了2个对象

            clean_img = crop_img(np.array(Image.open(clean_id).convert('RGB')), base=16)clean_patch_1, clean_patch_2 = self.crop_transform(clean_img), self.crop_transform(clean_img)# clean_name = self.clean_ids[clean_id].split("/")[-1].split('.')[0]clean_name = clean_id.split("/")[-1].split('.')[0]clean_patch_1, clean_patch_2 = random_augmentation(clean_patch_1, clean_patch_2)

crop_transform的定义如下,可见是随机进行crop

crop_transform = Compose([ToPILImage(),RandomCrop(args.patch_size),])

random_augmentation的实现代码如下,可以看到只是随机对图像进行翻转或旋转,其目的是尽可能使随机crop得到clean_patch_1, clean_patch_2差异更大,避免裁剪出高度相似的patch。

def random_augmentation(*args):out = []flag_aug = random.randint(1, 7)for data in args:out.append(data_augmentation(data, flag_aug).copy())return out
def data_augmentation(image, mode):if mode == 0:# originalout = image.numpy()elif mode == 1:# flip up and downout = np.flipud(image)elif mode == 2:# rotate counterwise 90 degreeout = np.rot90(image)elif mode == 3:# rotate 90 degree and flip up and downout = np.rot90(image)out = np.flipud(out)elif mode == 4:# rotate 180 degreeout = np.rot90(image, k=2)elif mode == 5:# rotate 180 degree and flipout = np.rot90(image, k=2)out = np.flipud(out)elif mode == 6:# rotate 270 degreeout = np.rot90(image, k=3)elif mode == 7:# rotate 270 degree and flipout = np.rot90(image, k=3)out = np.flipud(out)else:raise Exception('Invalid choice of image transformation')return out

3.2 degrad_patch分析

degrad_patch来自于clean_patch,可以看到是通过D.degrade进行转换的。

degrad_patch_1, degrad_patch_2 = self.D.degrade(clean_patch_1, clean_patch_2, de_id)

D.degrade相关的代码如下,可以看到只是对图像添加噪声。难怪AIRNet在图像去噪上效果最好。

class Degradation(object):def __init__(self, args):super(Degradation, self).__init__()self.args = argsself.toTensor = ToTensor()self.crop_transform = Compose([ToPILImage(),RandomCrop(args.patch_size),])def _add_gaussian_noise(self, clean_patch, sigma):# noise = torch.randn(*(clean_patch.shape))# clean_patch = self.toTensor(clean_patch)noise = np.random.randn(*clean_patch.shape)noisy_patch = np.clip(clean_patch + noise * sigma, 0, 255).astype(np.uint8)# noisy_patch = torch.clamp(clean_patch + noise * sigma, 0, 255).type(torch.int32)return noisy_patch, clean_patchdef _degrade_by_type(self, clean_patch, degrade_type):if degrade_type == 0:# denoise sigma=15degraded_patch, clean_patch = self._add_gaussian_noise(clean_patch, sigma=15)elif degrade_type == 1:# denoise sigma=25degraded_patch, clean_patch = self._add_gaussian_noise(clean_patch, sigma=25)elif degrade_type == 2:# denoise sigma=50degraded_patch, clean_patch = self._add_gaussian_noise(clean_patch, sigma=50)return degraded_patch, clean_patchdef degrade(self, clean_patch_1, clean_patch_2, degrade_type=None):if degrade_type == None:degrade_type = random.randint(0, 3)else:degrade_type = degrade_typedegrad_patch_1, _ = self._degrade_by_type(clean_patch_1, degrade_type)degrad_patch_2, _ = self._degrade_by_type(clean_patch_2, degrade_type)return degrad_patch_1, degrad_patch_2

4、推理演示

项目中默认包含了All.pth,要单独任务的模型可以到预训练模型下载地址: Google Drive and Baidu Netdisk (password: cr7d). 下载模型放到 ckpt/ 目录下

打开demo.py,将 subprocess.check_output(['mkdir', '-p', opt.output_path]) 替换为os.makedirs(opt.output_path,exist_ok=True),避免在window上报错,具体修改如下所示
在这里插入图片描述

demo.py默认从test\demo目录下读取图片进行测试,可见原始图像如下
在这里插入图片描述
代码运行后的输出结果默认保存在 output\demo目录下,可见对于去雨,去雾,去噪声效果都比较好。
在这里插入图片描述
模型推理时间如下所示,可以看到对一张320, 480的图片,要0.54s
在这里插入图片描述

这篇关于AIRNet模型使用与代码分析(All-In-One Image Restoration Network)的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1047130

相关文章

大模型研发全揭秘:客服工单数据标注的完整攻略

在人工智能(AI)领域,数据标注是模型训练过程中至关重要的一步。无论你是新手还是有经验的从业者,掌握数据标注的技术细节和常见问题的解决方案都能为你的AI项目增添不少价值。在电信运营商的客服系统中,工单数据是客户问题和解决方案的重要记录。通过对这些工单数据进行有效标注,不仅能够帮助提升客服自动化系统的智能化水平,还能优化客户服务流程,提高客户满意度。本文将详细介绍如何在电信运营商客服工单的背景下进行

中文分词jieba库的使用与实景应用(一)

知识星球:https://articles.zsxq.com/id_fxvgc803qmr2.html 目录 一.定义: 精确模式(默认模式): 全模式: 搜索引擎模式: paddle 模式(基于深度学习的分词模式): 二 自定义词典 三.文本解析   调整词出现的频率 四. 关键词提取 A. 基于TF-IDF算法的关键词提取 B. 基于TextRank算法的关键词提取

使用SecondaryNameNode恢复NameNode的数据

1)需求: NameNode进程挂了并且存储的数据也丢失了,如何恢复NameNode 此种方式恢复的数据可能存在小部分数据的丢失。 2)故障模拟 (1)kill -9 NameNode进程 [lytfly@hadoop102 current]$ kill -9 19886 (2)删除NameNode存储的数据(/opt/module/hadoop-3.1.4/data/tmp/dfs/na

Hadoop数据压缩使用介绍

一、压缩原则 (1)运算密集型的Job,少用压缩 (2)IO密集型的Job,多用压缩 二、压缩算法比较 三、压缩位置选择 四、压缩参数配置 1)为了支持多种压缩/解压缩算法,Hadoop引入了编码/解码器 2)要在Hadoop中启用压缩,可以配置如下参数

Makefile简明使用教程

文章目录 规则makefile文件的基本语法:加在命令前的特殊符号:.PHONY伪目标: Makefilev1 直观写法v2 加上中间过程v3 伪目标v4 变量 make 选项-f-n-C Make 是一种流行的构建工具,常用于将源代码转换成可执行文件或者其他形式的输出文件(如库文件、文档等)。Make 可以自动化地执行编译、链接等一系列操作。 规则 makefile文件

性能分析之MySQL索引实战案例

文章目录 一、前言二、准备三、MySQL索引优化四、MySQL 索引知识回顾五、总结 一、前言 在上一讲性能工具之 JProfiler 简单登录案例分析实战中已经发现SQL没有建立索引问题,本文将一起从代码层去分析为什么没有建立索引? 开源ERP项目地址:https://gitee.com/jishenghua/JSH_ERP 二、准备 打开IDEA找到登录请求资源路径位置

使用opencv优化图片(画面变清晰)

文章目录 需求影响照片清晰度的因素 实现降噪测试代码 锐化空间锐化Unsharp Masking频率域锐化对比测试 对比度增强常用算法对比测试 需求 对图像进行优化,使其看起来更清晰,同时保持尺寸不变,通常涉及到图像处理技术如锐化、降噪、对比度增强等 影响照片清晰度的因素 影响照片清晰度的因素有很多,主要可以从以下几个方面来分析 1. 拍摄设备 相机传感器:相机传

Andrej Karpathy最新采访:认知核心模型10亿参数就够了,AI会打破教育不公的僵局

夕小瑶科技说 原创  作者 | 海野 AI圈子的红人,AI大神Andrej Karpathy,曾是OpenAI联合创始人之一,特斯拉AI总监。上一次的动态是官宣创办一家名为 Eureka Labs 的人工智能+教育公司 ,宣布将长期致力于AI原生教育。 近日,Andrej Karpathy接受了No Priors(投资博客)的采访,与硅谷知名投资人 Sara Guo 和 Elad G

活用c4d官方开发文档查询代码

当你问AI助手比如豆包,如何用python禁止掉xpresso标签时候,它会提示到 这时候要用到两个东西。https://developers.maxon.net/论坛搜索和开发文档 比如这里我就在官方找到正确的id描述 然后我就把参数标签换过来

pdfmake生成pdf的使用

实际项目中有时会有根据填写的表单数据或者其他格式的数据,将数据自动填充到pdf文件中根据固定模板生成pdf文件的需求 文章目录 利用pdfmake生成pdf文件1.下载安装pdfmake第三方包2.封装生成pdf文件的共用配置3.生成pdf文件的文件模板内容4.调用方法生成pdf 利用pdfmake生成pdf文件 1.下载安装pdfmake第三方包 npm i pdfma