深度学习：基于MindSpore实现CycleGAN壁画修复

原创

已于 2024-10-05 01:02:52 修改 · 2.2k 阅读

30 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能

于 2024-10-05 01:02:22 首次发布

关于CycleGAN的基础知识可参考：

深度学习：CycleGAN图像风格迁移转换-优快云博客

以及MindSpore官方的教学视频：

CycleGAN图像风格迁移转换_哔哩哔哩_bilibili

本案例将基于CycleGAN实现破损草图到线稿图的转换

数据集

本案例使用的数据集里面的图片为经图线稿图数据。图像被统一缩放为256×256像素大小，其中用于训练的线稿图片25654张、草图图片25654张，用于测试的线稿图片100张、草图图片116张。

这里对数据进行了随机裁剪、水平随机翻转和归一化的预处理。

DexiNed

DexiNed（Dense Extreme Inception Network for Edge Detection）是一个为边缘检测任务设计的深度卷积神经网络模型。它由两个主要部分组成：Dexi和上采样网络（USNet）。

Dexi

这是模型的主要部分，包含六个编码块，每个块由多个子块组成，子块中包含卷积层、批量归一化和ReLU激活函数。从第二个块开始，引入了跳跃连接（skip connections），以保留不同层次的边缘特征。这些块的输出特征图被送入上采样网络以生成中间边缘图。

上采样网络（USNet）

这个部分由多个上采样块组成，每个块包括卷积层和反卷积层（也称为转置卷积层）。USNet的作用是将Dexi输出的低分辨率特征图上采样到更高的分辨率，以生成清晰的边缘图。

卷积层用于提取特征，反卷积层（或转置卷积层）用于将特征图的空间尺寸增大。

损失函数

DexiNed模型使用的损失函数是专门为边缘检测任务设计的，它在一定程度上受到了BDCN（Bi-directional Cascade Network）损失函数的启发，并进行了一些修改和优化。这个损失函数的目的是在训练过程中平衡正面（正样本）和负面（负样本）的边缘样本比例，从而提高边缘检测的准确性。

损失函数定义为：。

其中，

Li 是第 i 个输出的损失，λi 是对应的权重，用于平衡正负样本的比例。具体的 Li 计算方式为：

DexiNed数据集

DexiNed模型的训练数据集主要是为边缘检测任务设计的高质量数据集。在论文中提到了两个主要的数据集：

BIPED (Barcelona Images for Perceptual Edge Detection)：这是一个特别为边缘检测设计的大规模数据集，包含详细的边缘标注信息。它由250张真实世界的图像组成，图像分辨率为1280×720像素，主要描绘城市环境场景。这些图像的边缘通过手动标注生成，以确保边缘检测的准确性。
MDBD (Multicue Dataset for Boundary Detection)：这是一个用于边界检测的数据集，也适用于边缘检测任务。它由100个高清图像组成，每个图像有多个参与者的标注，适用于训练和评估边缘检测算法。

DexiNed模型需要成对的数据来进行训练，即每张输入图像都需要有一个对应的标注图像（Ground Truth, GT）。这些标注图像详细地标出了图像中边缘的位置，模型通过比较预测边缘和这些标注来学习如何准确地检测边缘。

DexiNed在本例中主要用于将彩色图片转化为线稿图，随后将线稿图输入CycleGAN，得到输出。

基于MindSpore的壁画修复

加载数据集

#下载数据集
from download import download

url = "https://6169fb4615b14dbcb6b2cb1c4eb78bb2.obs.cn-north-4.myhuaweicloud.com/Cyc_line.zip"

download(url, "./localdata", kind="zip", replace=True)

from __future__ import division
import math
import numpy as np

import os
import multiprocessing

import mindspore.dataset as de
import mindspore.dataset.vision as vision

"""数据集分布式采样器"""
class DistributedSampler:
    """Distributed sampler."""
    def __init__(self, dataset_size, num_replicas=None, rank=None, shuffle=True):
        if num_replicas is None:
            print("***********Setting world_size to 1 since it is not passed in ******************")
            num_replicas = 1
        if rank is None:
            print("***********Setting rank to 0 since it is not passed in ******************")
            rank = 0
        self.dataset_size = dataset_size
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.num_samples = int(math.ceil(dataset_size * 1.0 / self.num_replicas))
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle

    def __iter__(self):
        # deterministically shuffle based on epoch
        if self.shuffle:
            indices = np.random.RandomState(seed=self.epoch).permutation(self.dataset_size)
            # np.array type. number from 0 to len(dataset_size)-1, used as index of dataset
            indices = indices.tolist()
            self.epoch += 1
            # change to list type
        else:
            indices = list(range(self.dataset_size))

        # add extra samples to make it evenly divisible
        indices += indices[:(self.total_size - len(indices))]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)

    def __len__(self):
        return self.num_samples

# 加载CycleGAN数据集
IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.tif', '.tiff']

# 判断当前文件是否为图片
def is_image_file(filename):
    return any(filename.lower().endswith(extension) for extension in IMG_EXTENSIONS)

# 定义一个函数用于从指定目录中创建数据集列表
def make_dataset(dir_path, max_dataset_size=float("inf")):
    # 初始化一个空列表用来存储图片路径
    images = []
    # 确保提供的dir_path是一个有效的目录
    assert os.path.isdir(dir_path), '%s is not a valid directory' % dir_path
    # 遍历目录下的所有文件，将图片的文件路径存入images列表
    for root, _, fnames in sorted(os.walk(dir_path)):
        for fname in fnames:
            if is_image_file(fname):
                path = os.path.join(root, fname)
                images.append(path)
    # 返回指定长度的图片路径列表
    return images[:min(max_data