当前位置：首页 > article >正文

开源图像超分ECBSR项目源码分析

article 2024/10/23 3:37:46

在这里插入图片描述

ECBSR项目源码

源码说明

.
├── LICENSE
├── README.md
├── configs
│   ├── ecbsr_x2_m4c16_prelu.yml
│   ├── ecbsr_x2_m4c8_prelu.yml
│   ├── ecbsr_x4_m4c16_prelu.yml
│   └── ecbsr_x4_m4c8_prelu.yml
├── convert.py
├── datas
│   ├── __init__.py
│   ├── benchmark.py
│   └── div2k.py
├── deploy
│   ├── README.md
│   ├── ai-benchmark
│   │   └── README.md
│   ├── mininet
│   │   ├── Makefile
│   │   ├── README.md
│   │   ├── activation.cpp
│   │   ├── activation.h
│   │   ├── config.h
│   │   ├── convolution.cpp
│   │   ├── convolution.h
│   │   ├── elementwise.cpp
│   │   ├── elementwise.h
│   │   ├── main
│   │   ├── main.cpp
│   │   ├── padding.cpp
│   │   ├── padding.h
│   │   ├── pixelshuffle.cpp
│   │   ├── pixelshuffle.h
│   │   ├── tensor.cpp
│   │   ├── tensor.h
│   │   └── test.h
│   ├── mnn
│   │   └── README.md
│   └── rknn
│       └── README.MD
├── experiments
├── legacy
│   ├── README.md
│   └── src
│       ├── model
│       │   ├── ecb.py
│       │   └── ecbsr.py
│       └── option.py
├── models
│   ├── __init__.py
│   ├── ecb.py
│   ├── ecbsr.py
│   ├── plainsr.py
│   └── tf
│       ├── __init__.py
│       └── plainsr.py
├── requirements.txt
├── scripts
│   └── README.md
├── train.py
└── utils.py

安装依赖包文件：requirments.txt

imageio==2.9.0 //用于读写图像数据的Python库
numpy==1.18.0 //是一个开源的Python科学计算库，它提供了一个强大的N维数组对象和相应的工具集，用于快速操作数组。
pytorch_msssim==0.2.1 //是一个用于计算多尺度结构相似性（MS-SSIM）和结构相似性（SSIM）指数的快速且可微分的工具，
                       它基于 PyTorch 框架实现，并提供了与 TensorFlow 和 scikit-image 相一致的结果。
tqdm==4.36.1 //是一个快速、可扩展的Python进度条库，它可以在长循环中添加一个进度条，以便用户实时查看迭代过程的进度。
torchvision==0.9.0 //是 PyTorch 的一个扩展库，专门用于处理图像和视频数据，广泛应用于计算机视觉领域。
                     它提供了数据集、预训练模型、图像转换工具等，极大地简化了图像数据的处理过程。
scikit_image==0.15.0 // 是一个基于 Python 的图像处理库，
                     它建立在 NumPy、SciPy 和其他图像处理库之上，提供了丰富的图像处理算法和工具。
torch==1.8.1 //是一个开源的机器学习库，广泛用于计算机视觉和自然语言处理等应用。
PyYAML==5.4.1 //是一个用于处理 YAML 文件的 Python 库，它允许你轻松地加载和转储 YAML 格式的数据。
skimage==0.0 //Scikit-image（简称 skimage）是一个基于Python的图像处理库，
               它建立在NumPy和SciPy的基础上，提供了丰富的图像处理功能，适合于快速开发和实验。

数据集：DIV2K、benchmark
- DIV2K：https://cv.snu.ac.kr/research/EDSR/DIV2K.tar
- benchmark：https://cv.snu.ac.kr/research/EDSR/benchmark.tar

源码分析

数据处理模块

datas/div2k.py
- 功能：定义一个 DIV2K 的类，继承自 PyTorch 的 data.Dataset。这个类用于加载和处理 DIV2K 数据集；
- 逻辑图：
datas/benchmark.py
- 功能：定义一个 Benchmark 的类，继承自 PyTorch 的 data.Dataset。这个类用于加载和处理 Benchmark 数据集；
- 逻辑：

在这里插入图片描述

实用函数与类

utils.py
- 功能：定义了一些用于图像超分辨率任务的实用函数和类，包括计算峰值信噪比（PSNR）和结构相似性指数（SSIM），将HWC格式的NumPy数组转换为CHW格式的PyTorch张量，记录实验日志，以及跟踪和统计不同数据集的性能指标。
- 逻辑：

训练

train.py
- 功能：完成 ECBSR 算法的训练过程。需要导入数据处理类 Benchmark、DIV2K，以及超分网络模块类 ECBSR 类。
- 逻辑：

模型

models/ecbsr.py
- 功能：定义了一个 ECBSR 的类，核心是使用 ECB 的类作为主要得构建块
- 逻辑：
models/ecb.py
- 功能：实现 ECBSR 算法的核心神经网络模块。
- 逻辑：
models/plainsr.py
- 功能：是一个简单的超分辨率网络，它没有使用任何残差连接或复杂的结构，如残差块或密集块。
- 逻辑：

转换

convert.py
- 功能：将训练好的 model 从 pt 格式转换为 onnx 格式；需要导入ECBSR 类和 PlainSR 类
- 逻辑：

其他

legacy 文件夹：旧的版本
deploy 文件夹：部署相关介绍

注意

该项目没有推理，可以根据根据需求写个 inference 过程处理图片或视频，比如在gitcode仓库里增加了相关代码。

总结

ECBSR 算法针对图像 Y 通道进行训练和推理，因此输入图像转化为 YUV，对 Y 进行训练；推理时，需要将低分辨率图像上采样到高分辨率，然后将超分好的 Y 通道替换到上采样好的高分辨率图像，作为最后的输出超分图像。
ECBSR训练网络专门为训练推理使用，核心网络 ECB 跟论文介绍一样，是由3x3、1x1-3x3、1x1-sobelx、1x1-sobely、1x1-laplacian卷积层构成；在重参数可以将其权重和偏置组合合并后赋值给常规 3x3 卷积层作为PlainSR 网络的权重和偏置。
推理验证也利用跟训练一样的网络结构ECBSR，不过要开启模型的 eval 属性。
PlainSR 网络专门为移动端应用，核心网络跟论文中介绍一样，多个 3x3 卷积层构成，它没有使用任何残差连接或复杂的结构。
转换成 onnx 时，将 ECBSR 网络训练好的模型权重和偏置赋值给简单网络 Plainsr的权重和偏置，然后利用 torch.onnx.export导出 onnx 格式模型。