PixelCNN

import numpy as np
import keras
from keras import layers
from keras import ops
from tqdm import tqdm

# 模型 / 数据参数
num_classes = 10
input_shape = (28, 28, 1)
n_residual_blocks = 5
# 数据，分为训练集和测试集
(x, _), (y, _) = keras.datasets.mnist.load_data()
# 将所有图像连接在一起
data = np.concatenate((x, y), axis=0)
# 将所有小于最大值 256 的 33% 的像素值四舍五入为 0
# 任何高于此值的都四舍五入为 1，以使所有值均为
# 0 或 1
data = np.where(data < (0.33 * 256), 0, 1)
data = data.astype(np.float32)

# 第一个层是 PixelCNN 层。这个层简单地
# 建立在 2D 卷积层之上，但包含掩码。
class PixelConvLayer(layers.Layer):
    def __init__(self, mask_type, **kwargs):
        super().__init__()
        self.mask_type = mask_type
        self.conv = layers.Conv2D(**kwargs)

    def build(self, input_shape):
        # 建立 conv2d 层以初始化内核变量
        self.conv.build(input_shape)
        # 使用初始化的内核创建掩码
        kernel_shape = ops.shape(self.conv.kernel)
        self.mask = np.zeros(shape=kernel_shape)
        self.mask[: kernel_shape[0] // 2, ...] = 1.0
        self.mask[kernel_shape[0] // 2, : kernel_shape[1] // 2, ...] = 1.0
        if self.mask_type == "B":
            self.mask[kernel_shape[0] // 2, kernel_shape[1] // 2, ...] = 1.0

    def call(self, inputs):
        self.conv.kernel.assign(self.conv.kernel * self.mask)
        return self.conv(inputs)


# 接下来，我们构建我们的残差块层。
# 这只是一个普通的残差块，但基于 PixelConvLayer。
class ResidualBlock(keras.layers.Layer):
    def __init__(self, filters, **kwargs):
        super().__init__(**kwargs)
        self.conv1 = keras.layers.Conv2D(
            filters=filters, kernel_size=1, activation="relu"
        )
        self.pixel_conv = PixelConvLayer(
            mask_type="B",
            filters=filters // 2,
            kernel_size=3,
            activation="relu",
            padding="same",
        )
        self.conv2 = keras.layers.Conv2D(
            filters=filters, kernel_size=1, activation="relu"
        )

    def call(self, inputs):
        x = self.conv1(inputs)
        x = self.pixel_conv(x)
        x = self.conv2(x)
        return keras.layers.add([inputs, x])

inputs = keras.Input(shape=input_shape, batch_size=128)
x = PixelConvLayer(
    mask_type="A", filters=128, kernel_size=7, activation="relu", padding="same"
)(inputs)

for _ in range(n_residual_blocks):
    x = ResidualBlock(filters=128)(x)

for _ in range(2):
    x = PixelConvLayer(
        mask_type="B",
        filters=128,
        kernel_size=1,
        strides=1,
        activation="relu",
        padding="valid",
    )(x)

out = keras.layers.Conv2D(
    filters=1, kernel_size=1, strides=1, activation="sigmoid", padding="valid"
)(x)

pixel_cnn = keras.Model(inputs, out)
adam = keras.optimizers.Adam(learning_rate=0.0005)
pixel_cnn.compile(optimizer=adam, loss="binary_crossentropy")

pixel_cnn.summary()
pixel_cnn.fit(
    x=data, y=data, batch_size=128, epochs=50, validation_split=0.1, verbose=2
)

from IPython.display import Image, display

# 创建一个空的像素数组。
batch = 4
pixels = np.zeros(shape=(batch,) + (pixel_cnn.input_shape)[1:])
batch, rows, cols, channels = pixels.shape

# 遍历像素，因为生成必须按像素逐个顺序进行。
for row in tqdm(range(rows)):
    for col in range(cols):
        for channel in range(channels):
            # 输入整个数组并获取下一个像素的像素值概率。
            probs = pixel_cnn.predict(pixels)[:, row, col, channel]
            # 使用概率选择像素值并将其附加到图像框架中。
            pixels[:, row, col, channel] = ops.ceil(
                probs - keras.random.uniform(probs.shape)
            )


def deprocess_image(x):
    # 将单通道黑白图像堆叠为rgb值。
    x = np.stack((x, x, x), 2)
    # 取消预处理
    x *= 255.0
    # 转换为uint8并裁剪到有效范围[0, 255]
    x = np.clip(x, 0, 255).astype("uint8")
    return x


# 遍历生成的图像并使用matplotlib绘制它们。
for i, pic in enumerate(pixels):
    keras.utils.save_img(
        "generated_image_{}.png".format(i), deprocess_image(np.squeeze(pic, -1))
    )

display(Image("generated_image_0.png"))
display(Image("generated_image_1.png"))
display(Image("generated_image_2.png"))
display(Image("generated_image_3.png"))