<commit_before>#! /usr/bin/python
# -*- coding: utf8 -*-


import tensorflow as tf
import tensorlayer as tl
import numpy as np

import time
import numbers
import random
import os
import re
import sys

import threading
# import Queue  # <-- donot work for py3
is_py2 = sys.version[0] == '2'
if is_py2:
    import Queue as queue
else:
    import queue as queue

from six.moves import range
import scipy
from scipy import linalg
import scipy.ndimage as ndi

from skimage import transform
from skimage import exposure
import skimage

# linalg https://docs.scipy.org/doc/scipy/reference/linalg.html
# ndimage https://docs.scipy.org/doc/scipy/reference/ndimage.html

## Threading
def threading_data(data=None, fn=None, **kwargs):
    """Return a batch of result by given data.
    Usually be used for data augmentation.

    Parameters
    -----------
    data : numpy array, file names and etc, see Examples below.
    fn : the function for data processing.
    more args : the args for fn, see Examples below.

    Examples
    --------
    - Single array
    >>> X --> [batch_size, row, col, 1] greyscale
    >>> results = threading_data(X, zoom, zoom_range=[0.5, 1], is_random=True)
    ... results --> [batch_size, row, col, channel]
    >>> tl.visualize.images2d(images=np.asarray(results), second=0.01, saveable=True, name='after', dtype=None)
    >>> tl.visualize.images2d(images=np.asarray(X), second=0.01, saveable=True, name='before', dtype=None)

    - List of array (e.g. functions with ``multi``)
    >>> X, Y --> [batch_size, row, col, 1]  greyscale
    >>> data = threading_data([_ for _ in zip(X, Y)], zoom_multi, zoom_range=[0.5, 1], is_random=True)
    ... data --> [batch_size, 2, row, col, 1]
    >>> X_, Y_ = data.transpose((1,0,2,3,4))
    ... X_, Y_ --> [batch_size, row, col, 1]
    >>> tl.visualize.images2d(images=np.asarray(X_), second=0.01, saveable=True, name='after', dtype=None)
    >>> tl.visualize.images2d(images=np.asarray(Y_), second=0.01, saveable=True, name='before', dtype=None)

    - Customized function for image segmentation
    >>> def distort_img(data):
    ...     x, y = data
    ...     x, y = flip_axis_multi([x, y], axis=0, is_random=True)
    ...     x, y = flip_axis_multi([x, y], axis=1, is_random=True)
    ...     x, y = crop_multi([x, y], 100, 100, is_random=True)
    ...     return x, y
    >>> X, Y --> [batch_size, row, col, channel]
    >>> data = threading_data([_ for _ in zip(X, Y)], distort_img)
    >>> X_, Y_ = data.transpose((1,0,2,3,4))

    References
    ----------
    - `python queue <https://pymotw.com/2/Queue/index.html#module-Queue>`_
    - `run with limited queue <http://effbot.org/librarybook/queue.htm>`_
    """
    ## plot function info
    # for name, value in kwargs.items():
    #     print('{0} = {1}'.format(name, value))
    # exit()
    # define function for threading
    def apply_fn(results, i, data, kwargs):
        results[i] = fn(data, **kwargs)

    ## start multi-threaded reading.
    results = [None] * len(data) ## preallocate result list
    threads = []
    for i in range(len(data)):
        t = threading.Thread(
                        name='threading_and_return',
                        target=apply_fn,
                        args=(results, i, data[i], kwargs)
                        )
        t.start()
        threads.append(t)

    ## <Milo> wait for all threads to complete
    for t in threads:
        t.join()

    return np.asarray(results)

    ## old implementation
    # define function for threading
    # def function(q, i, data, kwargs):
    #     result = fn(data, **kwargs)
    #     q.put([i, result])
    # ## start threading
    # q = queue.Queue()
    # threads = []
    # for i in range(len(data)):
    #     t = threading.Thread(
    #                     name='threading_and_return',
    #                     target=function,
    #                     args=(q, i, data[i], kwargs)
    #                     )
    #     t.start()
    #     threads.append(t)
    #
    # ## <Milo> wait for all threads to complete
    # for t in threads:
    #     t.join()
    #
    # ## get results
    # results = []
    # for i in range(len(data)):
    #     result = q.get()
    #     results.append(result)
    # results = sorted(results)
    # for i in range(len(results)):
    #     results[i] = results[i][1]
    # return np.asarray(results)


## Image
def rotation(x, rg=20, is_random=False, row_index=0, col_index=1, channel_index=2,
                    fill_mode='nearest', cval=0.):
    """Rotate an image randomly or non-randomly.

    Parameters
    -----------
    x : numpy array
        An image with dimension of [row, col, channel] (default).
    rg : int or float
        Degree to rotate, usually 0 ~ 180.<commit_msg>Updated prepro.threading_data to use a certain number of threads. The function will now split the images into thread_count batches and assign each batch to a thread. Best results were using about 8 threads for 3x - 4x speed up (batch = 128) to 7x -8x (batch = 512) on tutorial_cifar10.py. Added prepro.pooling_data to use multiprocessing.Pool. There is an issue on Windows in that the pool must be created under __main__ somewhere (not in top level functions). Same performance as the new threading_data on 128 image batches in tutorial_cifar10.py, maybe 25% increase as batch size was increased. On my own data set, increase in performance of about 4x on new threading_data, and 16x on old threading_data for pool of 16 with 2048  64 x 64 images. All experiments on Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz with 16 cores and Titan X (pascal) GPU<commit_after>#! /usr/bin/python
# -*- coding: utf8 -*-


import tensorflow as tf
import tensorlayer as tl
import numpy as np

import time
import numbers
import random
import os
import re
import sys

import threading
# import Queue  # <-- donot work for py3
is_py2 = sys.version[0] == '2'
if is_py2:
    import Queue as queue
else:
    import queue as queue

from six.moves import range
import scipy
from scipy import linalg
import scipy.ndimage as ndi

from skimage import transform
from skimage import exposure
import skimage

from multiprocessing import Pool

# linalg https://docs.scipy.org/doc/scipy/reference/linalg.html
# ndimage https://docs.scipy.org/doc/scipy/reference/ndimage.html

## Threading
def threading_data(data=None, thread_count=None, fn=None, **kwargs):
    """Return a batch of result by given data.
    Usually be used for data augmentation.

    Parameters
    -----------
    data : numpy array, file names and etc, see Examples below.
    thread_count : the number of threads to use
    fn : the function for data processing.
    more args : the args for fn, see Examples below.

    Examples
    --------
    - Single array
    >>> X --> [batch_size, row, col, 1] greyscale
    >>> results = threading_data(X, None, zoom, zoom_range=[0.5, 1], is_random=True)
    ... results --> [batch_size, row, col, channel]
    >>> tl.visualize.images2d(images=np.asarray(results), second=0.01, saveable=True, name='after', dtype=None)
    >>> tl.visualize.images2d(images=np.asarray(X), second=0.01, saveable=True, name='before', dtype=None)

    - List of array (e.g. functions with ``multi``)
    >>> X, Y --> [batch_size, row, col, 1]  greyscale
    >>> data = threading_data([_ for _ in zip(X, Y)], None, zoom_multi, zoom_range=[0.5, 1], is_random=True)
    ... data --> [batch_size, 2, row, col, 1]
    >>> X_, Y_ = data.transpose((1,0,2,3,4))
    ... X_, Y_ --> [batch_size, row, col, 1]
    >>> tl.visualize.images2d(images=np.asarray(X_), second=0.01, saveable=True, name='after', dtype=None)
    >>> tl.visualize.images2d(images=np.asarray(Y_), second=0.01, saveable=True, name='before', dtype=None)
    
    - Single array split across ``thread_count`` threads (e.g. functions with ``multi``)
    >>> X, Y --> [batch_size, row, col, 1]  greyscale
    >>> data = threading_data(X, 8, zoom_multi, zoom_range=[0.5, 1], is_random=True)
    ... data --> [batch_size, 2, row, col, 1]
    >>> X_, Y_ = data.transpose((1,0,2,3,4))
    ... X_, Y_ --> [batch_size, row, col, 1]
    >>> tl.visualize.images2d(images=np.asarray(X_), second=0.01, saveable=True, name='after', dtype=None)
    >>> tl.visualize.images2d(images=np.asarray(Y_), second=0.01, saveable=True, name='before', dtype=None)

    - Customized function for image segmentation
    >>> def distort_img(data):
    ...     x, y = data
    ...     x, y = flip_axis_multi([x, y], axis=0, is_random=True)
    ...     x, y = flip_axis_multi([x, y], axis=1, is_random=True)
    ...     x, y = crop_multi([x, y], 100, 100, is_random=True)
    ...     return x, y
    >>> X, Y --> [batch_size, row, col, channel]
    >>> data = threading_data([_ for _ in zip(X, Y)], distort_img)
    >>> X_, Y_ = data.transpose((1,0,2,3,4))

    References
    ----------
    - `python queue <https://pymotw.com/2/Queue/index.html#module-Queue>`_
    - `run with limited queue <http://effbot.org/librarybook/queue.htm>`_
    """
    ## plot function info
    # for name, value in kwargs.items():
    #     print('{0} = {1}'.format(name, value))
    # exit()
    # define function for threading
    def apply_fn(results, i, data, kwargs):
        results[i] = fn(data, **kwargs)

    ## start multi-threaded reading.
    if thread_count is None:
        results = [None] * len(data) ## preallocate result list
        threads = []
        for i in range(len(data)):
            t = threading.Thread(
                            name='threading_and_return',
                            target=apply_fn,
                            args=(results, i, data[i], kwargs)
                            )
            t.start()
            threads.append(t)
    else:
        divs = np.linspace(0, len(data), thread_count + 1)
        divs = np.round(divs).astype(int)
        results = [None] * thread_count
        threads = []
        for i in range(thread_count):
            t = threading.Thread(
                name='threading_and_return',
                target=apply_fn,
                args=(results, i, data[divs[i]:divs[i + 1]], kwargs)
            )
            t.start()
            threads.append(t)

    ## <Milo> wait for all threads to complete
    for t in threads:
        t.join()

    if thread_count is None:
        return np.asarray(results)
    else:
        return np.concatenate(results)

## Pool
def pooling_data(data, pool: Pool, fn=None):
    """Return a batch of result by given data using multiprocessor.Pool
    Usually be used for data augmentation.

    Parameters
    -----------
    data : numpy array, file names and etc, see Examples below.
    pool : multiprocessing.Pool, the pool must be started from a function under ``if __name__ == "__main__":`` on Windows 
    fn : the function for data processing, must be a top level function with only one parameter

    Examples
    --------
    def distort_fn(x):
        x = tl.prepro.zoom(x, zoom_range=(0.5, 1.0), is_random=True, fill_mode='constant', cval=0.0)
        return x
        
    def main():
        # Setup graph ...
        # ...
        
        # Train
        pool = Pool(16)
        for i in range(epochs):       
            X_train_distorted = t.prepro.pooling_data(X_train, pool, distort_fn)
            # Batches
            for X, y in tl.iterate.minibatches(X_train_distorted, onehot_train, batch_size, shuffle=True):
                # ... rest of training code

    if __name__ == "__main__":
        main()
    
    References
    ----------
    - `reason for running Pool under __main__ <https://stackoverflow.com/questions/42602584/how-to-use-multiprocessing-pool-in-an-imported-module>`_
    """
    results = pool.map(fn, data)
    return np.asarray(results)

## Image
def rotation(x, rg=20, is_random=False, row_index=0, col_index=1, channel_index=2,
                    fill_mode='nearest', cval=0.):
    """Rotate an image randomly or non-randomly.

    Parameters
    -----------
    x : numpy array
        An image with dimension of [row, col, channel] (default).
    rg : int or float
        Degree to rotate, usually 0 ~ 180.
{'fraction_non_alphanumeric': 0.13964448119057463, 'fraction_numerical': 0.017610582885489872, 'mean_word_length': 2.4878892733564015, 'pattern_counts': {'":': 2, '<': 13, '<?xml version=': 0, '>': 119, 'https://': 7, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1641600', 'n_tokens_mistral': 4086, 'n_tokens_neox': 3914, 'n_words': 1147}