import pickle
import numpy as np
import os
import matplotlib.pyplot as plt

def load_pkls(path):
    assert os.path.isfile(path), path
    images = []
    with open(path, "rb") as f:
        images += pickle.load(f)
    assert len(images) > 0, path
    images = [np.transpose(image, [2, 0, 1]) for image in images]
    return images

path = 'datasets/DIV2K-va.pklv4'
loaded_images = load_pkls(path)
print(len(loaded_images))
# Display the first image
if loaded_images:
    first_image = loaded_images[11]
    plt.imshow(np.transpose(first_image, [1, 2, 0]))  # Transpose image to original shape [height, width, channels]
    plt.title('First Image')
    plt.axis('off')  # Hide axis
    plt.show()
else:
    print("No images loaded from the pickle file.")
print(loaded_images[11])