如何正确的将Mnist数据集(idx格式)解析成python数组？

import numpy as np
import struct

def parse_idx_file(filename):
    with open(filename, 'rb') as f:
        # 读取魔数和数据集长度
        magic_number = struct.unpack('>I', f.read(4))[0]
        data_length = struct.unpack('>I', f.read(4))[0]

        # 判断文件类型（图像或标签）
        if magic_number == 0x00000803:
            # 图像数据集
            num_rows = struct.unpack('>I', f.read(4))[0]
            num_cols = struct.unpack('>I', f.read(4))[0]

            # 读取像素数据
            images = np.frombuffer(f.read(), dtype=np.uint8).reshape(data_length, num_rows, num_cols)
            return images
        elif magic_number == 0x00000801:
            # 标签数据集
            labels = np.frombuffer(f.read(), dtype=np.uint8)
            return labels
        else:
            # 未知文件类型
            return None

train_images = parse_idx_file('train-images.idx3-ubyte')  # 训练集图像
train_labels = parse_idx_file('train-labels.idx1-ubyte')  # 训练集标签
test_images = parse_idx_file('t10k-images.idx3-ubyte')  # 测试集图像
test_labels = parse_idx_file('t10k-labels.idx1-ubyte')  # 测试集标签