dataset.py

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import os
import cv2
import h5py
import numpy as np
import tensorflow as tf

from glob import glob
from tqdm import tqdm
from util import split
from config import get_config
from multiprocessing import Pool


# Configuration
config, _ = get_config()

seed = config.seed


class DataSetLoader:

    @staticmethod
    def get_extension(ext):
        if ext in ['jpg', 'png']:
            return 'img'
        elif ext == 'tfr':
            return 'tfr'
        elif ext == 'h5':
            return 'h5'
        elif ext == 'npy':
            return 'npy'
        else:
            raise ValueError("[-] There's no supporting file... [%s] :(" % ext)

    @staticmethod
    def get_img(path, size=(64, 64), interp=cv2.INTER_CUBIC):
        img = cv2.imread(path, cv2.IMREAD_COLOR)[..., ::-1]  # BGR to RGB
        if img.shape[:1] == size:
            return img
        else:
            return cv2.resize(img, size, interp)

    @staticmethod
    def parse_tfr_tf(record):
        features = tf.parse_single_example(record, features={
            'shape': tf.FixedLenFeature([3], tf.int64),
            'data': tf.FixedLenFeature([], tf.string)})
        data = tf.decode_raw(features['data'], tf.uint8)
        return tf.reshape(data, features['shape'])

    @staticmethod
    def parse_tfr_np(record):
        ex = tf.train.Example()
        ex.ParseFromString(record)
        shape = ex.features.feature['shape'].int64_list.value
        data = ex.features.feature['data'].bytes_list.value[0]
        return np.fromstring(data, np.uint8).reshape(shape)

    @staticmethod
    def img_scaling(img, scale='0,1'):
        if scale == '0,1':
            try:
                img /= 255.
            except TypeError:  # ufunc 'true divide' output ~
                img = np.true_divide(img, 255.0, casting='unsafe')
        elif scale == '-1,1':
            try:
                img = (img / 127.5) - 1.
            except TypeError:
                img = np.true_divide(img, 127.5, casting='unsafe') - 1.
        else:
            raise ValueError("[-] Only '0,1' or '-1,1' please - (%s)" % scale)

        return img

    def __init__(self, path, size=None, name='to_tfr', use_save=False, save_file_name='',
                 buffer_size=4096, n_threads=8,
                 use_image_scaling=False, image_scale='0,1', img_save_method=cv2.INTER_LINEAR, debug=True):

        self.op = name.split('_')
        self.debug = debug

        try:
            assert len(self.op) == 2
        except AssertionError:
            raise AssertionError("[-] Invalid Target Types :(")

        self.size = size

        try:
            assert self.size
        except AssertionError:
            raise AssertionError("[-] Invalid Target Sizes :(")

        # To-DO
        # Supporting 4D Image
        self.height = size[0]
        self.width = size[1]
        self.channel = size[2]

        self.path = path

        try:
            assert os.path.exists(self.path)
        except AssertionError:
            raise AssertionError("[-] Path(%s) does not exist :(" % self.path)

        self.buffer_size = buffer_size
        self.n_threads = n_threads

        if os.path.isfile(self.path):
            self.file_list = [self.path]
            self.file_ext = self.path.split('.')[-1]
            self.file_names = [self.path]
        else:
            self.file_list = sorted(os.listdir(self.path))
            self.file_ext = self.file_list[0].split('.')[-1]
            self.file_names = glob(self.path + '/*')
        self.raw_data = np.ndarray([], dtype=np.uint8)  # (N, H * W * C)

        if self.debug:
            print("[*] Detected Path            is [%s]" % self.path)
            print("[*] Detected File Extension  is [%s]" % self.file_ext)
            print("[*] Detected First File Name is [%s] (%d File(s))" % (self.file_names[0], len(self.file_names)))

        self.types = ('img', 'tfr', 'h5', 'npy')  # Supporting Data Types
        self.op_src = self.get_extension(self.file_ext)
        self.op_dst = self.op[1]

        try:
            chk_src, chk_dst = False, False
            for t in self.types:
                if self.op_src == t:
                    chk_src = True
                if self.op_dst == t:
                    chk_dst = True
            assert chk_src and chk_dst
        except AssertionError:
            raise AssertionError("[-] Invalid Operation Types (%s, %s) :(" % (self.op_src, self.op_dst))

        self.img_save_method = img_save_method

        if self.op_src == self.types[0]:
            self.load_img()
        elif self.op_src == self.types[1]:
            self.load_tfr()
        elif self.op_src == self.types[2]:
            self.load_h5()
        elif self.op_src == self.types[3]:
            self.load_npy()
        else:
            raise NotImplementedError("[-] Not Supported Type :(")

        # Random Shuffle
        order = np.arange(self.raw_data.shape[0])
        np.random.RandomState(seed).shuffle(order)
        self.raw_data = self.raw_data[order]

        # Clip [0, 255]
        self.raw_data = np.rint(self.raw_data).clip(0, 255).astype(np.uint8)

        self.use_save = use_save
        self.save_file_name = save_file_name

        if self.use_save:
            try:
                assert self.save_file_name
            except AssertionError:
                raise AssertionError("[-] Empty save-file name :(")

            if self.op_dst == self.types[0]:
                self.convert_to_img()
            elif self.op_dst == self.types[1]:
                self.tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE)
                self.tfr_writer = tf.python_io.TFRecordWriter(self.save_file_name + ".tfrecords", self.tfr_opt)
                self.convert_to_tfr()
            elif self.op_dst == self.types[2]:
                self.convert_to_h5()
            elif self.op_dst == self.types[3]:
                self.convert_to_npy()
            else:
                raise NotImplementedError("[-] Not Supported Type :(")

        self.use_image_scaling = use_image_scaling
        self.img_scale = image_scale

        if self.use_image_scaling:
            self.raw_data = self.img_scaling(self.raw_data, self.img_scale)

    def load_img(self):
        self.raw_data = np.zeros((len(self.file_list), self.height * self.width * self.channel),
                                 dtype=np.uint8)

        for i, fn in tqdm(enumerate(self.file_names)):
            self.raw_data[i] = self.get_img(fn, (self.height, self.width), self.img_save_method).flatten()
            if self.debug:  # just once
                print("[*] Image Shape   : ", self.raw_data[i].shape)
                print("[*] Image Size    : ", self.raw_data[i].size)
                print("[*] Image MIN/MAX :  (%d, %d)" % (np.min(self.raw_data[i]), np.max(self.raw_data[i])))
                self.debug = False

    def load_tfr(self):
        self.raw_data = tf.data.TFRecordDataset(self.file_names, compression_type='', buffer_size=self.buffer_size)
        self.raw_data = self.raw_data.map(self.parse_tfr_tf, num_parallel_calls=self.n_threads)

    def load_h5(self, size=0, offset=0):
        init = True

        for fl in self.file_list:  # For multiple .h5 files
            with h5py.File(fl, 'r') as hf:
                data = hf['images']
                full_size = len(data)

                if size == 0:
                    size = full_size

                n_chunks = int(np.ceil(full_size / size))
                if offset >= n_chunks:
                    print("[*] Looping from back to start.")
                    offset %= n_chunks
                if offset == n_chunks - 1:
                    print("[-] Not enough data available, clipping to end.")
                    data = data[offset * size:]
                else:
                    data = data[offset * size:(offset + 1) * size]

                data = np.array(data, dtype=np.uint8)
                print("[+] ", fl, " => Image size : ", data.shape)

                if init:
                    self.raw_data = data
                    init = False

                    if self.debug:  # just once
                        print("[*] Image Shape   : ", self.raw_data[0].shape)
                        print("[*] Image Size    : ", self.raw_data[0].size)
                        print("[*] Image MIN/MAX :  (%d, %d)" % (np.min(self.raw_data[0]), np.max(self.raw_data[0])))
                        self.debug = False

                    continue
                else:
                    self.raw_data = np.concatenate((self.raw_data, data))

    def load_npy(self):
        self.raw_data = np.rollaxis(np.squeeze(np.load(self.file_names), axis=0), 0, 3)

        if self.debug:  # just once
            print("[*] Image Shape   : ", self.raw_data[0].shape)
            print("[*] Image Size    : ", self.raw_data[0].size)
            print("[*] Image MIN/MAX :  (%d, %d)" % (np.min(self.raw_data[0]), np.max(self.raw_data[0])))
            self.debug = False

    def convert_to_img(self):
        def to_img(i):
            cv2.imwrite('imgHQ%05d.png' % i, cv2.COLOR_BGR2RGB)
            return True

        raw_data_shape = self.raw_data.shape  # (N, H * W * C)

        try:
            assert os.path.exists(self.save_file_name)
        except AssertionError:
            print("[-] There's no %s :(" % self.save_file_name)
            print("[*] Make directory at %s... " % self.save_file_name)
            os.mkdir(self.save_file_name)

        ii = [i for i in range(raw_data_shape[0])]

        pool = Pool(self.n_threads)
        print(pool.map(to_img, ii))

    def convert_to_tfr(self):
        for data in self.raw_data:
            ex = tf.train.Example(features=tf.train.Features(feature={
                'shape': tf.train.Feature(int64_list=tf.train.Int64List(value=data.shape)),
                'data': tf.train.Feature(bytes_list=tf.train.BytesList(value=[data.tostring()]))
            }))
            self.tfr_writer.write(ex.SerializeToString())

    def convert_to_h5(self):
        with h5py.File(self.save_file_name, 'w') as f:
            f.create_dataset("images", data=self.raw_data)

    def convert_to_npy(self):
        np.save(self.save_file_name, self.raw_data)


class Div2KDataSet:

    def __init__(self, hr_height=768, hr_width=768, lr_height=192, lr_width=192, channel=3,
                 n_patch=16, use_split=False, split_rate=0.1, random_state=42, n_threads=8,
                 ds_path=None, ds_name=None, use_img_scale=True,
                 ds_hr_path=None, ds_lr_path=None,
                 use_save=False, save_type='to_h5', save_file_name=None, debug=False):

        """
        # General Settings
        :param hr_height: input HR image height, default 768
        :param hr_width: input HR image width, default 768
        :param lr_height: input LR image height, default 192
        :param lr_width: input LR image width, default 192
        :param channel: input image channel, default 3 (RGB)
        - in case of Div2K - ds x4, image size is 768 x 768 x 3 (HWC).

        # Pre-Processing Option
        :param n_patch: patch size to crop, default 16
        :param split_rate: image split rate (into train & test), default 0.1
        :param random_state: random seed for shuffling, default 42
        :param n_threads: the number of threads for multi-threading, default 8

        # DataSet Option
        :param ds_path: DataSet's Path, default None
        :param ds_name: DataSet's Name, default None
        :param use_img_scale: using img scaling?, default False
        :param ds_hr_path: DataSet High Resolution path
        :param ds_lr_path: DataSet Low Resolution path
        :param use_save: saving into another file format
        :param save_type: file format to save
        :param save_file_name: file name to save
        :param debug: debugging messages, default False
        """

        self.hr_height = hr_height
        self.hr_width = hr_width
        self.lr_height = lr_height
        self.lr_width = lr_width
        self.channel = channel
        self.hr_shape = (self.hr_height, self.hr_width, self.channel)
        self.lr_shape = (self.lr_height, self.lr_width, self.channel)

        self.n_patch = n_patch
        self.use_split = use_split
        self.split_rate = split_rate
        self.random_state = random_state
        self.num_threads = n_threads  # change this value to the fitted value for ur system

        """
        Expected ds_path : div2k/...
        Expected ds_name : X4
        """
        self.ds_path = ds_path
        self.ds_name = ds_name
        self.ds_hr_path = ds_hr_path
        self.ds_lr_path = ds_lr_path

        try:
            assert self.ds_path
        except AssertionError:
            try:
                assert self.ds_hr_path and self.ds_lr_path
            except AssertionError:
                raise AssertionError("[-] DataSet's path is required!")

        self.use_save = use_save
        self.save_type = save_type
        self.save_file_name = save_file_name
        self.debug = debug

        try:
            if self.use_save:
                assert self.save_file_name
            else:
                self.save_file_name = ""
        except AssertionError:
            raise AssertionError("[-] save-file/folder-name is required!")

        self.n_images = 800
        self.n_images_val = 100

        self.use_img_scaling = use_img_scale

        if self.ds_path:  # like .h5 or .tfr # will be in the same folder
            self.ds_hr_path = self.ds_path + "/DIV2K_train_HR/"
            self.ds_lr_path = self.ds_hr_path

        self.hr_images = DataSetLoader(path=self.ds_hr_path,
                                       size=self.hr_shape,
                                       use_save=self.use_save,
                                       name=self.save_type,
                                       save_file_name=self.save_file_name + "-hr.h5",
                                       use_image_scaling=self.use_img_scaling,
                                       image_scale='0,1',
                                       img_save_method=cv2.INTER_LINEAR).raw_data  # numpy arrays
        self.patch_hr_images = None

        self.lr_images = DataSetLoader(path=self.ds_lr_path,
                                       size=self.lr_shape,
                                       use_save=self.use_save,
                                       name=self.save_type,
                                       save_file_name=self.save_file_name + "-lr.h5",
                                       use_image_scaling=self.use_img_scaling,
                                       image_scale='0,1',
                                       img_save_method=cv2.INTER_CUBIC).raw_data  # numpy arrays
        self.patch_lr_images = None

        if self.n_patch > 0:
            patch_size = int(np.sqrt(self.n_patch))

            self.patch_hr_images = np.zeros((self.n_images * self.n_patch,
                                             self.hr_height // patch_size, self.hr_width // patch_size, self.channel),
                                            dtype=np.uint8)

            self.patch_lr_images = np.zeros((self.n_images * self.n_patch,
                                             self.lr_height // patch_size, self.lr_width // patch_size, self.channel),
                                            dtype=np.uint8)

            for i in tqdm(range(self.n_images)):
                hr_patches = split(np.reshape(self.hr_images[i, :], self.hr_shape), self.n_patch)
                lr_patches = split(np.reshape(self.lr_images[i, :], self.lr_shape), self.n_patch)

                for n_ps in range(self.n_patch):
                    self.patch_hr_images[i * self.n_patch + n_ps] = hr_patches[n_ps]
                    self.patch_lr_images[i * self.n_patch + n_ps] = lr_patches[n_ps]

                if self.debug:
                    import matplotlib.pyplot as plt

                    fig = plt.figure()
                    for j in range(self.n_patch):
                        fig.add_subplot(4, 4, j + 1)
                        plt.imshow(self.patch_hr_images[j, :, :, :])
                    plt.show()

                    fig = plt.figure()
                    for j in range(self.n_patch):
                        fig.add_subplot(4, 4, j + 1)
                        plt.imshow(self.patch_lr_images[j, :, :, :])
                    plt.show()

                    self.debug = False


class DataIterator:

    def __init__(self, x, y, batch_size):
        self.x = x
        self.y = y
        self.batch_size = batch_size
        self.num_examples = num_examples = x.shape[0]
        self.num_batches = num_examples // batch_size
        self.pointer = 0

        assert (self.batch_size <= self.num_examples)

    def next_batch(self):
        start = self.pointer
        self.pointer += self.batch_size

        if self.pointer > self.num_examples:
            perm = np.arange(self.num_examples)
            np.random.shuffle(perm)

            self.x = self.x[perm, :, :, :]
            self.y = self.y[perm, :, :, :]

            start = 0
            self.pointer = self.batch_size

        end = self.pointer
        return self.x[start:end, :, :, :], self.y[start:end, :, :, :]

    def iterate(self):
        for step in range(self.num_batches):
            yield self.next_batch()