如何编辑mnist数据集？

import numpy as np
import cv2
from tensorflow.keras.datasets import mnist

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 定义图像变换函数
def rotate_image(image, angle):
    return cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)

# 对训练数据进行旋转增强
augmented_images = []
augmented_labels = []
for image, label in zip(x_train, y_train):
    augmented_images.append(image)
    augmented_labels.append(label)
    rotated_image = rotate_image(image, 90)
    augmented_images.append(rotated_image)
    augmented_labels.append(label)

# 转换为numpy数组
augmented_images = np.array(augmented_images)
augmented_labels = np.array(augmented_labels)

# 打印增强后的数据集大小
print("Augmented images shape:", augmented_images.shape)
print("Augmented labels shape:", augmented_labels.shape)