prepare_train.py

import numpy as np
import cv2 as cv
import os
import scipy.misc
import h5py

# Settings.
scale = 3
size_input = 33
size_label = 21
stride = 14
counter = 0
# data = np.zeros([size_input, size_input, 1, 1])
# label = np.zeros([size_label, size_label, 1, 1])
data = []
label = []
padding = abs(size_input - size_label) / 2

# Source: https://stackoverflow.com/questions/29100722/equivalent-im2double-function-in-opencv-python
def im2double(im):
    info = np.iinfo(im.dtype) # Get the data type of the input image
    return im.astype(np.float) / info.max # Divide all values by the largest possible value in the datatype

# Make sure the image shape is divisible by the scale factor. Also convert 3-channel image to 1-channel image.
def modcrop(image, scale=3):
	if image.shape[2] == 1:
		size = image.shape
		size -= np.mod(size, scale)
		image = image[0:size[0], 0:size[1]]
	else:
		size = image.shape[0:2]
		size -= np.mod(size, scale)
		image = image[0:size[0], 0:size[1], 0]
	return image

# Load and preprocess the training images.
dirpath = './Train/'
for root, dirs, files in os.walk(dirpath):
	for file in files:
		# Read in image and convert to ycrcb color space.
		img = cv.imread(dirpath + file)
		# cv.imshow('image',img)
		# cv.waitKey(0)
		# cv.destroyAllWindows()
		img = cv.cvtColor(img, cv.COLOR_BGR2YCR_CB)
		img = im2double(img) # Only use the luminance value.

		# Create groundtruth and baseline image.
		im_label = modcrop(img)
		size = im_label.shape
		h = size[0]
		w = size[1]
		im_temp = scipy.misc.imresize(im_label, 1/scale, interp='bicubic')
		im_input = scipy.misc.imresize(im_temp, size, interp='bicubic')

		# Generate subimages for training.
		for x in range(0, h - size_input, stride):
			for y in range(0, w - size_input, stride):
				subim_input = im_input[x : x + size_input, y : y + size_input]
				subim_label = im_label[int(x + padding) : int(x + padding + size_label), int(y + padding) : int(y + padding + size_label)]
				
				subim_input = subim_input.reshape([size_input, size_input, 1])
				subim_label = subim_label.reshape([size_label, size_label, 1])

				data.append(subim_input)
				label.append(subim_label)
				counter += 1

# Shuffle the data pairs.
order = np.random.choice(counter, counter, replace=False)
data = np.array([data[i] for i in order])
label = np.array([label[i] for i in order])

print('data shape is', data.shape)
print('label shape is', label.shape)

# Write to HDF5 file.
savepath = os.path.join(os.getcwd(), 'checkpoint/train.h5')
with h5py.File(savepath, 'w') as hf:
    hf.create_dataset('data', data=data)
    hf.create_dataset('label', data=label)