from datasets import load_dataset
import numpy as np
import matplotlib.pyplot as plt

def main():
    dataset = load_dataset('parquet', data_files={
        'train': r"D:\AiData\Dataset\mnist\mnist\train-00000-of-00001.parquet",
        'test': r"D:\AiData\Dataset\mnist\mnist\test-00000-of-00001.parquet",
    })
    train_dataset = dataset['train']
    first_sample = train_dataset[0]

    print("Label:", first_sample['label'])

    image = first_sample['image']
    image_array = np.array(image)
    print("Image shape:", image_array.shape)

    plt.imshow(image_array, cmap='gray')
    plt.show()

    print("First few rows of pixel values:")
    print(image_array[:5, :5])


if __name__ == "__main__":
    main()