import glob import h5py from PIL import Image from torchvision.transforms import RandomCrop from torchvision.transforms.functional import to_tensor from tqdm import tqdm from Dataloader import ImageAugment patch_size = 128 shrink_size = 2 noise_level = 1 patches_per_img = 20 images = glob.glob("dataset/train/*") database = h5py.File("train_images.hdf5", 'w') dat_group = database.create_group("shrink_2_noise_level_1_downsample_random_rgb") # del database['shrink_2_noise_level_1_downsample_random'] storage_lr = dat_group.create_dataset("train_lr", shape=(patches_per_img * len(images), 3, patch_size // shrink_size, patch_size // shrink_size), dtype='float32', # compression='lzf', ) storage_hr = dat_group.create_dataset("train_hr", shape=(patches_per_img * len(images), 3, patch_size, patch_size), # compression='lzf', dtype='float32') random_cropper = RandomCrop(size=patch_size) img_augmenter = ImageAugment(shrink_size, noise_level, down_sample_method=None) def get_img_patches(img_pil): img_patch = random_cropper(img_pil) lr_hr_patches = img_augmenter.process(img_patch) return lr_hr_patches counter = 0 for img in tqdm(images): img_pil = Image.open(img).convert("RGB") for i in range(patches_per_img): patch = get_img_patches(img_pil) storage_lr[counter] = to_tensor(patch[0].convert("RGB")).numpy() storage_hr[counter] = to_tensor(patch[1].convert("RGB")).numpy() counter += 1 database.close()