AetherNet: Image Restoration with Design Patterns

特性

设计模式

1. 注册表模式 (Registry Pattern)

from core import MODEL_REGISTRY
@MODEL_REGISTRY.register()
class MyModel(nn.Module):
 ...
# 获取已注册模型
model_cls = MODEL_REGISTRY.get('MyModel')

2. 工厂模式 (Factory Pattern)

from core import ModelFactory
# 通过配置创建
config = {'model_name': 'SwinIR', 'upscale': 4, 'embed_dim': 180}
model = ModelFactory.create(config)
# 通过名称创建
model = ModelFactory.create_by_name('CWFNet', upscale=4)

3. 策略模式 (Strategy Pattern)

from core import ClassicalSRDegradation, ColorDenoiseDegradation
# 经典SR退化(双三次下采样)
degradation = ClassicalSRDegradation(scale=4)
lr_image = degradation.apply(hr_image)
# 去噪退化(高斯噪声)
degradation = ColorDenoiseDegradation(sigma=25)
noisy_image = degradation.apply(clean_image)

4. 模板方法模式 (Template Method Pattern)

from core import BaseTrainer
class SwinIRTrainer(BaseTrainer):
 def build_model(self):
 return SwinIR(**self.config['model'])
 
 def train_step(self, batch):
 lr, hr = batch
 output = self.model(lr)
 loss = self.criterion(output, hr)
 ...
 return loss.item()
trainer = SwinIRTrainer(config)
trainer.train()

安装

环境要求

快速安装

git clone https://github.com/Jackksonns/AetherNet.git
cd AetherNet
# 创建虚拟环境(推荐)
conda create -n aethernet python=3.9
conda activate aethernet
# 安装PyTorch(根据CUDA版本选择)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
# 安装依赖
pip install -r requirements.txt
# MambaIRv2支持(可选)
pip install mamba-ssm

数据集准备

DIV2K

# 下载DIV2K数据集
cd datasets
bash organize_div2k.sh
# 目录结构
datasets/
├── DIV2K/
│ ├── train/
│ │ ├── HR/ # 高分辨率图像
│ │ └── LR_bicubic/ # 双三次下采样图像
│ │ ├── X2/
│ │ ├── X3/
│ │ └── X4/
│ └── val/

自定义数据集

使用方法

训练

# 使用配置文件训练
python scripts/train/train.py --config configs/train_MambaIRv2_SR_x4.yml
# 训练SwinIR
python scripts/train/train_swinir.py --scale 4 --batch_size 8 --epochs 100
# 训练CWFNet
python scripts/train/train_cwfnet.py --scale 4 --use_wavelet --use_cvoca

测试

# 测试SwinIR
python scripts/test/test_swinir.py --checkpoint weights/swinir_x4.pth
# 测试CWFNet
python scripts/test/test_cwfnet.py --checkpoint weights/cwfnet_x4.pth

Python API

from models import SwinIR
from core import ModelFactory
import torch
# 方式1:直接导入
model = SwinIR(upscale=4, in_chans=3, img_size=64, 
 window_size=8, img_range=1., depths=[6]*6,
 embed_dim=180, num_heads=[6]*6, mlp_ratio=2,
 upsampler='pixelshuffle', resi_connection='1conv')
# 方式2:工厂创建
model = ModelFactory.create_by_name('SwinIR', upscale=4)
# 加载权重
model.load_state_dict(torch.load('weights/model.pth'))
model.eval()
# 推理
with torch.no_grad():
 sr_image = model(lr_image)

实验结果

经典图像超分辨率 (×ばつ)

配置说明

YAML配置文件示例

# configs/train_swinir_x4.yml
model:
 name: SwinIR
 upscale: 4
 in_chans: 3
 img_size: 64
 window_size: 8
 embed_dim: 180
 depths: [6, 6, 6, 6, 6, 6]
 num_heads: [6, 6, 6, 6, 6, 6]
 mlp_ratio: 2
 upsampler: pixelshuffle
training:
 batch_size: 8
 epochs: 100
 lr: 1e-4
 scheduler:
 type: cosine
 
data:
 train_dir: datasets/DIV2K/train
 val_dir: datasets/DIV2K/val
 patch_size: 64

引用

致谢