大纲

1.介绍
2.代码模块

介绍

图像超分辨率复原（Super-Resolution Restoration, SR）是指采用图像处理和机器学习技术，从同一场景中已有的一张或多张低分辨率图像构建高分辨率图像的技术。在图像获取过程中，由于成像距离、成像设备分辨率等因素的限制，成像系统难以无失真地获取原始场景中的信息。按照傅里叶光学的观点，光学成像系统是一个低通滤波器，由于受到光学衍射的影响，其传递函数在由衍射极限分辨率所决定的某个截止频率以上值均为零。超分辨率图像复原的目的就在于复原截止频率之外的信息，以使图像获得更多的细节和信息。

本次我所采用的模型为carn模型，主打轻量化，参数量仅40万左右，但是他的结果数据还不错。
夜景图的PSNR为四十五点一六分贝，SSIM为零点九九三二。头发图的PSNR为四十五点九零分贝，SSIM为零点九九三八。

代码模块

包含carn.py，test.py，train.py。训练一次，然后跑test，注意改成自己的路径。

carn.py

"""
CARN: Convolutional Anchored Regression Network
Paper: https://arxiv.org/abs/1803.08664
PyTorch Implementation
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class MeanShift(nn.Module):
    """数据标准化层"""
    def __init__(self, mean_rgb, sub):
        super(MeanShift, self).__init__()
        sign = -1 if sub else 1
        r, g, b = mean_rgb
        self.shift = nn.Parameter(torch.tensor([r, g, b]) * sign, requires_grad=False)
        self.register_buffer('rgb_range', torch.tensor(1.0))

    def forward(self, x):
        return x + self.shift.view(1, 3, 1, 1)


class ResBlock(nn.Module):
    """基础残差块"""
    def __init__(self, n_feats, kernel_size=3, bias=True, bn=False, act=nn.ReLU(True)):
        super(ResBlock, self).__init__()
        m = []
        for i in range(2):
            m.append(nn.Conv2d(n_feats, n_feats, kernel_size, padding=kernel_size//2, bias=bias))
            if bn:
                m.append(nn.BatchNorm2d(n_feats))
            if i == 0:
                m.append(act)
        self.body = nn.Sequential(*m)

    def forward(self, x):
        return x + self.body(x)


class Upsampler(nn.Module):
    """上采样模块"""
    def __init__(self, n_feats, scale, bn=False, act=False, bias=True):
        super(Upsampler, self).__init__()
        m = []
        if (scale & (scale - 1)) == 0:  # scale 是 2 的幂
            for _ in range(int(np.log2(scale))):
                m.append(nn.Conv2d(n_feats, 4 * n_feats, 3, padding=1, bias=bias))
                m.append(nn.PixelShuffle(2))
                if bn:
                    m.append(nn.BatchNorm2d(n_feats))
                if act:
                    m.append(act())
        elif scale == 3:
            m.append(nn.Conv2d(n_feats, 9 * n_feats, 3, padding=1, bias=bias))
            m.append(nn.PixelShuffle(3))
            if bn:
                m.append(nn.BatchNorm2d(n_feats))
            if act:
                m.append(act())
        else:
            raise NotImplementedError
        self.body = nn.Sequential(*m)

    def forward(self, x):
        return self.body(x)


class Block(nn.Module):
    """CARN 核心回归块"""
    def __init__(self, n_feats, kernel_size=3, reduction=16):
        super(Block, self).__init__()
        self.res1 = ResBlock(n_feats, kernel_size, bias=True, bn=False)
        self.res2 = ResBlock(n_feats, kernel_size, bias=True, bn=False)
        self.res3 = ResBlock(n_feats, kernel_size, bias=True, bn=False)

        # 级联连接：1x1 卷积融合
        self.c1 = nn.Conv2d(n_feats * 2, n_feats, 1, bias=True)
        self.c2 = nn.Conv2d(n_feats * 3, n_feats, 1, bias=True)
        self.c3 = nn.Conv2d(n_feats * 4, n_feats, 1, bias=True)

    def forward(self, x):
        res1 = self.res1(x)
        concat1 = torch.cat([x, res1], dim=1)
        f1 = self.c1(concat1)

        res2 = self.res2(f1)
        concat2 = torch.cat([x, res1, res2], dim=1)
        f2 = self.c2(concat2)

        res3 = self.res3(f2)
        concat3 = torch.cat([x, res1, res2, res3], dim=1)
        f3 = self.c3(concat3)

        return f3


class CARN(nn.Module):
    """
    CARN: Convolutional Anchored Regression Network
    Args:
        scale: 放大倍数 (2, 3, 4)
        n_colors: 输入通道数 (RGB=3)
        n_feats: 特征通道数 (默认 64)
        n_blocks: 回归块数量 (默认 3)
        reduction: 通道缩减因子
    """
    def __init__(self, scale=4, n_colors=3, n_feats=64, n_blocks=3, reduction=16):
        super(CARN, self).__init__()

        self.n_colors = n_colors
        self.n_feats = n_feats
        self.scale = scale

        # RGB 均值归一化
        rgb_mean = (0.4488, 0.4371, 0.4040)
        self.sub_mean = MeanShift(rgb_mean, sub=True)
        self.add_mean = MeanShift(rgb_mean, sub=False)

        # 头部：特征提取
        self.head = nn.Conv2d(n_colors, n_feats, 3, padding=1, bias=True)

        # 主体：级联回归块
        blocks = []
        for _ in range(n_blocks):
            blocks.append(Block(n_feats, reduction=reduction))
        self.body = nn.Sequential(*blocks)

        # 尾部：上采样 + 重建
        self.upsampler = Upsampler(n_feats, scale)
        self.tail = nn.Conv2d(n_feats, n_colors, 3, padding=1, bias=True)

    def forward(self, x):
        # 归一化
        x = self.sub_mean(x)

        # 特征提取
        x = self.head(x)

        # 主体
        res = self.body(x)
        res = res + x  # 全局残差连接

        # 上采样 + 重建
        res = self.upsampler(res)
        x = self.tail(res)

        # 反归一化
        x = self.add_mean(x)

        return x


class CARN_M(nn.Module):
    """
    CARN-Mobile: 轻量版 CARN，使用组卷积减少参数量
    参数量约 0.41M，显存占用更低，适合 8GB 显卡
    """
    def __init__(self, scale=4, n_colors=3, n_feats=64, n_blocks=3, reduction=4):
        super(CARN_M, self).__init__()

        self.n_colors = n_colors
        self.n_feats = n_feats
        self.scale = scale

        rgb_mean = (0.4488, 0.4371, 0.4040)
        self.sub_mean = MeanShift(rgb_mean, sub=True)
        self.add_mean = MeanShift(rgb_mean, sub=False)

        # 使用组卷积 (groups=n_feats) 实现深度可分离卷积
        self.head = nn.Conv2d(n_colors, n_feats, 3, padding=1, bias=True)

        blocks = []
        for _ in range(n_blocks):
            blocks.append(Block(n_feats, reduction=reduction))
        self.body = nn.Sequential(*blocks)

        self.upsampler = Upsampler(n_feats, scale)
        self.tail = nn.Conv2d(n_feats, n_colors, 3, padding=1, bias=True)

    def forward(self, x):
        x = self.sub_mean(x)
        x = self.head(x)
        res = self.body(x)
        res = res + x
        res = self.upsampler(res)
        x = self.tail(res)
        x = self.add_mean(x)
        return x


# 为了 Upsampler 中使用 np
import numpy as np

test.py

"""
CARN 测试脚本：对比输入低分辨率图、Bicubic插值放大、CARN超分结果，并显示PSNR/SSIM
热力图采用 2%~98% 分位数裁剪，增强对比度，清晰显示边缘差异
"""

import os
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 禁用 cuDNN 自动优化，减少显存分配
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

# 中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Zen Hei']
plt.rcParams['axes.unicode_minus'] = False

from carn import CARN_M, CARN

try:
    from skimage.metrics import peak_signal_noise_ratio, structural_similarity
except ImportError:
    raise ImportError("请安装 scikit-image：pip install scikit-image -i https://pypi.tuna.tsinghua.edu.cn/simple")


def load_model(model_path, model_type="carn_m", scale=4, device="cuda"):
    if model_type == "carn_m":
        model = CARN_M(scale=scale)
    else:
        model = CARN(scale=scale)

    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()

    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    return model


def sr_image(model, lr_path, output_path=None, lr_output_path=None, device="cuda", max_size=256):
    """
    读取图片，缩放到 max_size，然后进行超分
    返回：sr_img (numpy, RGB), lr_original (缩放后的低分辨率图, numpy, RGB)
    参数：
        lr_path: 原始图片路径
        output_path: 超分结果保存路径（可选）
        lr_output_path: 缩放后的低分辨率图保存路径（可选）
    """
    img = cv2.imread(lr_path)
    if img is None:
        raise FileNotFoundError(f"图片不存在：{lr_path}")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # 缩放
    h, w = img.shape[:2]
    if max(h, w) > max_size:
        scale = max_size / max(h, w)
        new_w = int(w * scale)
        new_h = int(h * scale)
        img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
        print(f"图片缩放至：{new_h}x{new_w}（原图 {h}x{w}）")
    lr_original = img.copy()

    # 保存缩放后的低分辨率图
    if lr_output_path:
        lr_bgr = cv2.cvtColor(lr_original, cv2.COLOR_RGB2BGR)
        cv2.imwrite(lr_output_path, lr_bgr)
        print(f"缩放后的低分辨率图已保存至：{lr_output_path}")

    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    # 转 tensor
    img_tensor = torch.from_numpy(img.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)

    with torch.no_grad():
        sr_tensor = model(img_tensor)

    sr_img = sr_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
    sr_img = np.clip(sr_img * 255, 0, 255).astype(np.uint8)

    del img_tensor, sr_tensor
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    if output_path:
        sr_img_bgr = cv2.cvtColor(sr_img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(output_path, sr_img_bgr)
        print(f"超分结果已保存至：{output_path}")

    return sr_img, lr_original


def compare_and_visualize(lr_original, sr_img, save_path=None):
    """
    对比：LR输入、Bicubic放大、SR输出，计算PSNR/SSIM
    显示4列：输入LR、Bicubic、CARN-SR、误差热力图（2%~98%裁剪）
    """
    h_sr, w_sr = sr_img.shape[:2]

    # Bicubic 放大 LR 到 SR 尺寸
    lr_bicubic = cv2.resize(lr_original, (w_sr, h_sr), interpolation=cv2.INTER_CUBIC)

    # 计算指标
    psnr = peak_signal_noise_ratio(lr_bicubic, sr_img, data_range=255)
    min_side = min(h_sr, w_sr)
    if min_side < 3:
        ssim = None
    else:
        win_size = 7 if min_side >= 7 else (3 if min_side >= 3 else min_side)
        if win_size % 2 == 0:
            win_size -= 1
        try:
            ssim = structural_similarity(lr_bicubic, sr_img, win_size=win_size, channel_axis=-1, data_range=255)
        except:
            ssim = None

    # 计算误差
    diff = np.abs(lr_bicubic.astype(np.float32) - sr_img.astype(np.float32))
    print(f"误差统计: 最小值={diff.min():.2f}, 最大值={diff.max():.2f}, 均值={diff.mean():.2f}, 标准差={diff.std():.2f}")

    # ----- 使用 2%~98% 分位数裁剪，增强对比度 -----
    vmin_val = np.percentile(diff, 2)
    vmax_val = np.percentile(diff, 98)
    if vmin_val >= vmax_val:  # 防止图片太均匀
        vmin_val = 0
        vmax_val = diff.max()
    print(f"热力图裁剪范围: vmin={vmin_val:.2f}, vmax={vmax_val:.2f}")

    # 可视化：4列
    fig, axes = plt.subplots(1, 4, figsize=(18, 5))

    # 1. LR 输入
    axes[0].imshow(lr_original)
    axes[0].set_title(f"输入图\n{lr_original.shape[0]}x{lr_original.shape[1]}", fontsize=12)
    axes[0].axis('off')

    # 2. Bicubic 放大
    axes[1].imshow(lr_bicubic)
    axes[1].set_title(f"Bicubic 上采样\n({lr_bicubic.shape[0]}x{lr_bicubic.shape[1]})", fontsize=12)
    axes[1].axis('off')

    # 3. CARN 超分结果
    axes[2].imshow(sr_img)
    ssim_str = f"SSIM: {ssim:.4f}" if ssim is not None else "SSIM: N/A"
    axes[2].set_title(f"CARN-M 超分结果\nPSNR: {psnr:.2f} dB, {ssim_str}", fontsize=12)
    axes[2].axis('off')

    # 4. 误差热力图（2%~98% 裁剪）
    im = axes[3].imshow(diff, cmap='hot', vmin=vmin_val, vmax=vmax_val)
    axes[3].set_title(f"误差热力图 (增强对比度)\n越亮差异越大", fontsize=12)
    axes[3].axis('off')
    plt.colorbar(im, ax=axes[3], fraction=0.046, pad=0.04)

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='t'
                                                    'ight')
        print(f"对比图已保存至：{save_path}")
    plt.show()

    return psnr, ssim


if __name__ == "__main__":
    MODEL_PATH = "./checkpoints/best_model.pth"
    MODEL_TYPE = "carn_m"
    SCALE = 2
    TEST_IMAGE = "./A.png"
    OUTPUT_SR = "./A_result.png"
    OUTPUT_LR = "./A_resized.png"
    OUTPUT_COMPARE = "./A_comparison.png"
    MAX_SIZE = 1024

    # 设备选择
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("使用设备: GPU")
    else:
        device = torch.device("cpu")
        print("未检测到 GPU，使用 CPU")

    # 加载模型
    model = load_model(MODEL_PATH, MODEL_TYPE, SCALE, device)

    # 执行超分，同时保存缩放后的低分辨率图
    sr_img, lr_img = sr_image(
        model,
        TEST_IMAGE,
        output_path=OUTPUT_SR,
        lr_output_path=OUTPUT_LR,
        device=device,
        max_size=MAX_SIZE
    )

    # 显示对比与指标
    psnr, ssim = compare_and_visualize(lr_img, sr_img, OUTPUT_COMPARE)

    print(f"\n✅ 定量评价结果：PSNR = {psnr:.2f} dB" + (f"，SSIM = {ssim:.4f}" if ssim else ""))

train.py

"""
CARN 训练脚本
针对 RTX 4070 8GB 显存优化: batch_size=4, 使用 CARN-M 轻量版
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Compose, ToTensor, RandomCrop, RandomHorizontalFlip, RandomVerticalFlip
import cv2
import numpy as np
from tqdm import tqdm
from loguru import logger
from carn import CARN_M, CARN

# ============ 配置 ============
class Config:
    # 数据
    train_dir = "./data/train"      # 训练集 HR 图片目录
    val_dir = "./data/val"          # 验证集 HR 图片目录
    scale = 2                       # 放大倍数 (2, 3, 4)
    patch_size = 48                 # 训练 patch 大小
    batch_size = 4                  # 8GB 显存推荐 4 (CARN-M 可尝试 8)
    num_workers = 2

    # 模型
    model_type = "carn_m"           # "carn" 或 "carn_m" (推荐 carn_m)
    n_feats = 64
    n_blocks = 3

    # 训练
    epochs = 100
    lr = 1e-4
    lr_step = 50
    lr_gamma = 0.5
    save_interval = 10

    # 保存
    save_dir = "./checkpoints"
    log_interval = 10


class SRDataset(Dataset):
    """超分数据集：从 HR 图片生成 LR-HR 训练对"""
    def __init__(self, img_dir, scale=2, patch_size=48, augment=True):
        self.img_dir = img_dir
        self.scale = scale
        self.patch_size = patch_size
        self.augment = augment

        # 读取所有图片路径
        self.img_paths = []
        for ext in ['.png', '.jpg', '.jpeg', '.bmp']:
            self.img_paths.extend([os.path.join(img_dir, f) for f in os.listdir(img_dir) if f.endswith(ext)])

        logger.info(f"Found {len(self.img_paths)} images in {img_dir}")

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        # 读取 HR 图片
        hr = cv2.imread(self.img_paths[idx])
        hr = cv2.cvtColor(hr, cv2.COLOR_BGR2RGB)
        hr = hr.astype(np.float32) / 255.0

        # 随机裁剪
        h, w = hr.shape[:2]
        ph = self.patch_size
        pw = self.patch_size

        # 确保裁剪区域有效 (HR 尺寸需 >= patch_size * scale)
        min_h = ph * self.scale
        min_w = pw * self.scale
        if h < min_h or w < min_w:
            hr = cv2.resize(hr, (min_w, min_h))

        h, w = hr.shape[:2]
        top = np.random.randint(0, h - ph * self.scale + 1)
        left = np.random.randint(0, w - pw * self.scale + 1)
        hr_patch = hr[top:top + ph * self.scale, left:left + pw * self.scale]

        # 生成 LR
        lr_patch = cv2.resize(hr_patch, (pw, ph), interpolation=cv2.INTER_CUBIC)

        # 数据增强
        if self.augment:
            if np.random.random() > 0.5:
                hr_patch = np.fliplr(hr_patch).copy()
                lr_patch = np.fliplr(lr_patch).copy()
            if np.random.random() > 0.5:
                hr_patch = np.flipud(hr_patch).copy()
                lr_patch = np.flipud(lr_patch).copy()

        # 转换为 Tensor: (C, H, W)
        hr_tensor = torch.from_numpy(hr_patch.transpose(2, 0, 1)).float()
        lr_tensor = torch.from_numpy(lr_patch.transpose(2, 0, 1)).float()

        return lr_tensor, hr_tensor


def train():
    # 设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")

    # 显存信息
    if torch.cuda.is_available():
        logger.info(f"GPU: {torch.cuda.get_device_name(0)}")
        logger.info(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

    # 数据加载
    train_dataset = SRDataset(Config.train_dir, Config.scale, Config.patch_size, augment=True)
    train_loader = DataLoader(train_dataset, batch_size=Config.batch_size,
                              shuffle=True, num_workers=Config.num_workers, pin_memory=True)

    # 模型
    if Config.model_type == "carn_m":
        model = CARN_M(scale=Config.scale, n_feats=Config.n_feats, n_blocks=Config.n_blocks)
        logger.info("Using CARN-M (lightweight)")
    else:
        model = CARN(scale=Config.scale, n_feats=Config.n_feats, n_blocks=Config.n_blocks)
        logger.info("Using CARN")

    model = model.to(device)

    # 打印参数量
    params = sum(p.numel() for p in model.parameters())
    logger.info(f"Total parameters: {params:,} ({params / 1e6:.2f}M)")

    # 损失函数 & 优化器
    criterion = nn.L1Loss()  # L1 loss 比 L2 收敛更好
    optimizer = optim.Adam(model.parameters(), lr=Config.lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=Config.lr_step, gamma=Config.lr_gamma)

    # 创建保存目录
    os.makedirs(Config.save_dir, exist_ok=True)

    # 训练循环
    logger.info("Start training...")
    best_loss = float('inf')

    for epoch in range(1, Config.epochs + 1):
        model.train()
        epoch_loss = 0.0

        pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{Config.epochs}")
        for i, (lr, hr) in enumerate(pbar):
            lr = lr.to(device)
            hr = hr.to(device)

            optimizer.zero_grad()
            sr = model(lr)
            loss = criterion(sr, hr)
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()
            pbar.set_postfix({"loss": f"{loss.item():.4f}"})

        avg_loss = epoch_loss / len(train_loader)

        # 学习率调整
        scheduler.step()

        # 打印日志
        logger.info(f"Epoch {epoch}: Avg Loss = {avg_loss:.6f}, LR = {scheduler.get_last_lr()[0]:.2e}")

        # 保存最佳模型
        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': best_loss,
            }, os.path.join(Config.save_dir, 'best_model.pth'))
            logger.info(f"Best model saved (loss: {best_loss:.6f})")

        # 定期保存
        if epoch % Config.save_interval == 0:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': avg_loss,
            }, os.path.join(Config.save_dir, f'checkpoint_epoch_{epoch}.pth'))

    logger.info(f"Training completed! Best loss: {best_loss:.6f}")


if __name__ == "__main__":
    train()