import h5py
from PIL import Image
import numpy as np

def save_multiple_he_images(hdf5_file_path):
    # Define the dataset names to loop over
    he_datasets = ['HE/raw', 'HE/sn1', 'HE/sn2']

    with h5py.File(hdf5_file_path, 'r') as file:
        for dataset_name in he_datasets:
            # Access the dataset
            dataset = file[dataset_name]
            # Iterate over each image in the dataset
            for i, img in enumerate(dataset):
                # Convert the numpy array to a PIL Image
                image = Image.fromarray(img.astype(np.uint8))
                # Construct a fitting filename for each image
                filename = f'{dataset_name.replace("/", "_")}_{i}.png'
                # Save the image
                image.save(filename)
                print(f"Saved {filename}")

# Example usage, replace with your actual HDF5 file path
hdf5_file_path = '/lustre/groups/shared/CoCaHis/file.hdf5'
save_multiple_he_images(hdf5_file_path)