生成式压缩

基于生成对抗网络的图像学习压缩 TensorFlow 实现。该方法由 Agustsson 等人在 极端学习图像压缩的生成对抗网络 中提出。他们的想法非常有趣，且方法描述得十分清晰。

---

使用方法

代码依赖于 TensorFlow 1.8

# 克隆
$ git clone https://github.com/Justin-Tan/generative-compression.git
$ cd generative-compression

# 训练时，请查看命令行参数
$ python3 train.py -h
# 运行
$ python3 train.py -opt momentum --name my_network

训练采用 batch size 1，重建样本和 TensorBoard 摘要会每隔一定步数（默认为 128 步）定期写入。每 10 个 epoch 保存一次检查点。

压缩单张图片：

# 压缩
$ python3 compress.py -r /path/to/model/checkpoint -i /path/to/image -o path/to/output/image

压缩后的图像将与原始图像并排对比，保存在 config.py 中 directories.samples 指定的路径下。如果您使用的是提供的预训练模型并进行噪声采样，请保留 config.py 中 config_test 下的超参数；否则，测试时的参数应与训练时设置的参数一致。

注： 如果您愿意以更高的比特率换取更高的感知质量，可以尝试实现“高保真生成式图像压缩”(High-Fidelity Generative Image Compression)，它与本项目思路相同，但运行在更高的比特率区间。此外，该实现还能处理任意大小和分辨率的图像。

结果

这些全局压缩的图像来自 Cityscapes leftImg8bit 数据集的测试集。解码器似乎会在建筑物中“幻觉”出绿色植被，反之亦然。

全局条件压缩：多尺度判别器 + 特征匹配损失，C=8 通道 - （压缩至 0.072 bpp）

第 38 个 epoch 第 44 个 epoch 第 47 个 epoch 第 48 个 epoch

展示量化后的 C=4、8、16 通道图像对比

生成器损失	判别器损失

预训练模型

您可以在下方找到全局压缩的预训练模型，其通道瓶颈为 C = 8（对应 0.072 bpp 表示）。该模型采用了多尺度判别器和特征匹配损失。噪声从 128 维正态分布中采样，经过类似 DCGAN 的生成器处理后与量化后的图像表示拼接。模型在 Cityscapes leftImg8bit 数据集的训练集上训练了 55 个 epoch，并使用 gtFine 数据集作为对应的语义标签图。此模型应能配合 config.py 中 config_test 的默认设置正常工作。

同时附带一个带有 C=8 瓶颈的全局条件压缩预训练模型。该模型同样使用上述损失函数训练了 50 个 epoch。重建过程以语义标签图为条件（详见 cGAN/ 文件夹及“条件 GAN 使用说明”）。

• 噪声采样模型
• 条件 GAN 模型

警告： 模型是用 TensorFlow 1.3 训练的，但在 TensorFlow 1.8 上加载并无问题。如有任何问题，请提交 issue。

详情/扩展

网络架构基于原始论文附录中的描述，而该附录又参考了论文 用于实时风格迁移和超分辨率的感知损失。所使用的多尺度判别器损失最初是在 使用条件 GAN 进行高分辨率图像合成与语义操作 项目中提出的，具体实现请参阅 network.py。若需添加扩展功能，可在 Network 类中创建新方法，例如：

@staticmethod
def my_generator(z, **kwargs):
    """
    输入：
    z: 采样的噪声

    输出：
    上采样后的图像
    """

    return tf.random_normal([z.get_shape()[0], height, width, channels], seed=42)

如需更改超参数或切换功能，可使用 config.py 中的配置项。（虽然可能不太规范，但我认为这样比编写 20 行的 argparse 更方便）。

数据/准备

训练使用了 ADE 20k 数据集 和 Cityscapes leftImg8bit 数据集。前者将图像缩放至宽度 512 像素，后者则在训练前将图像重新采样至 [512 x 1024]（参见 Imagemagick 命令行选项）。data/ 目录下提供了使用 Imagemagick 重新采样的示例脚本。在每种情况下，您都需要创建一个包含单列 path 的 Pandas 数据框，其中存储图像的绝对或相对路径。该数据框应保存为 HDF5 文件，并将其路径提供给 config.py 中的 directories 类。data 目录中提供了 Cityscapes 数据集的示例。

条件 GAN 使用

全局压缩的条件 GAN 实现位于 cGAN 目录中。条件 GAN 实现似乎能生成质量最高的图像，但目前仍处于实验阶段。在此实现中，生成过程以所选图像的语义标签图信息为条件。您需要下载 gtFine 注释地图数据集，并在 Pandas 数据框中添加一列 semantic_map_paths，指向 gtFine 数据集中对应的图像。

依赖项

• Python 3.6
• Pandas
• TensorFlow 1.8

待办事项：

• 将 GAN 噪声采样融入重建图像中。作者指出这一步是可选的，采样的噪声会与量化表示相结合，但未提供更多细节。目前模型从正态分布中采样，并使用类似 DCGAN 的生成器上采样（参见 network.py），以将其与量化图像表示 w_hat 拼接在一起，但这似乎显著增加了重建图像中的“幻觉”成分。
• 集成 VGG 损失。
• 尝试 WGAN-GP。
• 尝试谱归一化。
• 探索不同带有噪声采样的生成器架构。
• 扩展到使用语义图进行选择性压缩（欢迎贡献）。

资源

• 用于极端学习图像压缩的生成对抗网络
• CycleGAN
• 基于条件 GAN 的高分辨率图像合成与语义操控

更多结果

全局压缩：噪声采样、多尺度判别器 + 特征匹配损失，C=8 通道 - 压缩至 0.072 bpp

generative-compression 快速上手指南

本项目是基于 TensorFlow 的生成对抗网络（GAN）图像压缩实现，源自论文《Generative Adversarial Networks for Extreme Learned Image Compression》。它能够在极低码率下实现具有良好感知质量的图像重建。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux / macOS (Windows 需自行配置 TensorFlow 1.x 环境)
• Python 版本: Python 3.6
• 核心框架: TensorFlow 1.8 (模型训练基于 TF 1.3，但在 TF 1.8 上加载正常)
• 其他依赖: Pandas

依赖安装命令：

pip install tensorflow==1.8 pandas

提示：由于 TensorFlow 1.x 版本较老，若国内下载缓慢，建议使用清华或阿里镜像源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow==1.8 pandas

安装步骤

1. 克隆仓库 将项目代码下载到本地：

   git clone https://github.com/Justin-Tan/generative-compression.git
   cd generative-compression
   

2. 数据准备（可选） 若需重新训练模型，需准备数据集（如 Cityscapes 或 ADE 20k）。

   • 创建一个包含单列 path 的 Pandas DataFrame，存储图像的绝对或相对路径。
   • 将其保存为 HDF5 文件。
   • 在 config.py 的 directories 类中指定该文件的路径。
   • 参考 data/ 目录下的示例脚本进行图像重采样处理。

基本使用

1. 训练模型

运行训练脚本前，可查看命令行参数帮助：

python3 train.py -h

启动训练示例（使用 momentum 优化器，命名网络为 my_network）：

python3 train.py -opt momentum --name my_network

• 说明：默认 Batch Size 为 1。每隔一定步数（默认 128 步）会写入重建样本和 TensorBoard 摘要，每 10 个 epoch 保存一次检查点。

2. 压缩单张图像

使用预训练模型或训练好的检查点对单张图片进行压缩：

python3 compress.py -r /path/to/model/checkpoint -i /path/to/image -o path/to/output/image

• -r: 模型检查点路径。
• -i: 输入图像路径。
• -o: 输出图像路径。

注意：

• 压缩后的图像将以“原图 - 重建图”并排对比的形式保存，具体保存位置由 config.py 中的 directories.samples 定义。
• 若使用提供的带噪声采样的预训练模型，请保持 config.py 中 config_test 下的超参数不变；若是自定义训练模型，测试时的参数需与训练时保持一致。

3. 使用预训练模型

项目提供了两种全局压缩预训练模型（瓶颈通道数 C=8，约 0.072 bpp）：

• 噪声采样模型 (Noise sampling model): 下载链接
• 条件 GAN 模型 (Conditional GAN model): 下载链接 (需配合语义标签图使用)

下载后解压，将路径填入上述 compress.py 命令的 -r 参数即可运行。