Search code examples
pythonopencvdatasethdf5h5py

Creating a custom image dataset for super-resolution

I'm working on developing a custom image dataset for a super-resolution deep learning network. I have the images saved to disk and can create the HDF5 dataset files. Here's the code I'm using:

import os, cv2, h5py, glob
import numpy as np
from glob import glob

# define the paths to the dataset
BASE_DATA_PATH = '/usr/local/home/.../esrgan_data'
HR_TRAIN_PATH = os.path.join(BASE_DATA_PATH, 'train_HR')
LR_TRAIN_PATH = os.path.join(BASE_DATA_PATH, 'train_LR')

# create LR and HR image lists
LR_images = glob(LR_TRAIN_PATH + '**/*.png')
HR_images = glob(HR_TRAIN_PATH + '**/*.png')

# sort the lists
LR_images.sort()
HR_images.sort()

# create an h5 file
with h5py.File('datasets/esrgan_trainDS.h5', 'w') as h5_file:
    # create 2 datasets for LR and HR images in the h5 file
    lr_ds = h5_file.create_dataset('trainLR', (len(LR_images), 150, 150, 3), dtype='f')
    hr_ds = h5_file.create_dataset('trainHR', (len(HR_images), 600, 600, 3), dtype='f')
    
    for i in range(len(LR_images)):
        LR_image = cv2.imread(LR_images[i])
        HR_image = cv2.imread(HR_images[i])
        lr_trainDS[i] = LR_image
        hr_trainDS[i] = HR_image

# load the h5 dataset
trainDS = h5py.File('datasets/esrgan_trainDS.h5', 'r')
print('Files in the training dataset: ', list(trainDS.keys()))

Files in the training dataset: ['trainHR', 'trainLR']

LRset = trainDS['trainLR']
HRset = trainDS['trainHR']

print('LR dataset shape: ', LRset.shape)
print('HR dataset shape: ', HRset.shape)

LR dataset shape: (450, 150, 150, 3) HR dataset shape: (450, 600, 600, 3)

My problem is that when I try to view an individual image from the dataset, I see a black box which tells me the image either didn't save or didn't load properly.

cv2_imshow('', HRset[100])

enter image description here

I based the code on this post. The code runs without error -- I can write the f5 files, read them and print file attributes. I just can't see the images and, without an error message, I'm not sure where I'm going wrong.

I'm guessing it's a simple mistake I'm not seeing, but I'd appreciate any help you're able to provide. Thanks!

Solution

  • You are the victim of a small error in the referenced post. (My apologies; I was the author of that post. Apparently no one noticed the error creating the dataset with the wrong dtype. I corrected that answer to get the dtype from an image, then use it when creating the dataset.)

    In addition, you need to add some lines to your code to view the image. Once you fix the HDF5 file, add these 2 lines to your code to view the image. (Also, correct the small typo on imshow()). Modified code below:

    cv2.imshow('', HRset[100]) 
cv2.waitKey(0) # waits until a key is pressed in the image window
cv2.destroyAllWindows() # destroys the window showing the image

    What do these lines do?
    cv2.waitKey(0) pauses the execution of your program. As a result, the image window will remain visible. If you do not include this statement, cv2.imshow() executes in the blink of the eye, then the program closes all windows it opened. That makes it very unlikely you will see the image in the window. (The black image you posted is the remnants of the image window.)