<s> from __future__ import division , print_function , unicode_literals <EOL> import math <EOL> import numpy as np <EOL> import six <EOL> from brainstorm . handlers . _cpuop import _crop_images <EOL> from brainstorm . randomness import Seedable <EOL> from brainstorm . utils import IteratorValidationError <EOL> class DataIterator ( Seedable ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , data_shapes , length ) : <EOL> """<STR_LIT>""" <EOL> super ( DataIterator , self ) . __init__ ( ) <EOL> self . data_shapes = data_shapes <EOL> self . length = length <EOL> def __call__ ( self , handler ) : <EOL> pass <EOL> class AddGaussianNoise ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , std_dict , mean_dict = None ) : <EOL> """<STR_LIT>""" <EOL> DataIterator . __init__ ( self , iter . data_shapes , iter . length ) <EOL> mean_keys = set ( mean_dict . keys ( ) ) if mean_dict is not None else set ( ) <EOL> std_keys = set ( std_dict . keys ( ) ) <EOL> if mean_dict is not None and mean_keys != std_keys : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( mean_keys , std_keys ) ) <EOL> for key in std_keys : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> self . mean_dict = { } if mean_dict is None else mean_dict <EOL> self . std_dict = std_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler = None ) : <EOL> for data in self . iter ( handler ) : <EOL> for key , std in self . std_dict . items ( ) : <EOL> mean = self . mean_dict . get ( key , <NUM_LIT:0.0> ) <EOL> data [ key ] = data [ key ] + std * self . rnd . standard_normal ( <EOL> data [ key ] . shape ) + mean <EOL> yield data <EOL> class AddSaltNPepper ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , prob_dict , ratio_dict = None ) : <EOL> """<STR_LIT>""" <EOL> DataIterator . __init__ ( self , iter . data_shapes , iter . length ) <EOL> ratio_keys = set ( ) if ratio_dict is None else set ( ratio_dict . keys ( ) ) <EOL> prob_keys = set ( prob_dict . keys ( ) ) <EOL> if ratio_dict is not None and ratio_keys != prob_keys : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( prob_keys , ratio_keys ) ) <EOL> for key in prob_keys : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> self . ratio_dict = { } if ratio_dict is None else ratio_dict <EOL> self . prob_dict = prob_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler = None ) : <EOL> for data in self . iter ( handler ) : <EOL> for key , pr in self . prob_dict . items ( ) : <EOL> ratio = self . ratio_dict . get ( key , <NUM_LIT:0.5> ) <EOL> d = data [ key ] . copy ( ) <EOL> r = self . rnd . rand ( * d . shape ) <EOL> d [ r >= <NUM_LIT:1.0> - pr * ratio ] = <NUM_LIT:1.0> <EOL> d [ r <= pr * ( <NUM_LIT:1.0> - ratio ) ] = <NUM_LIT:0.0> <EOL> data [ key ] = d <EOL> yield data <EOL> class Flip ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , prob_dict = None ) : <EOL> """<STR_LIT>""" <EOL> Seedable . __init__ ( self ) <EOL> super ( Flip , self ) . __init__ ( iter . data_shapes , iter . length ) <EOL> prob_dict = { '<STR_LIT:default>' : <NUM_LIT:0.5> } if prob_dict is None else prob_dict <EOL> for key in prob_dict . keys ( ) : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> if prob_dict [ key ] > <NUM_LIT:1.0> or prob_dict [ key ] < <NUM_LIT:0.0> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> if len ( iter . data_shapes [ key ] ) != <NUM_LIT:5> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> self . prob_dict = prob_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler = None ) : <EOL> for data in self . iter ( handler ) : <EOL> for name in self . prob_dict . keys ( ) : <EOL> assert isinstance ( data [ name ] , np . ndarray ) <EOL> for i in range ( data [ name ] . shape [ <NUM_LIT:1> ] ) : <EOL> if self . rnd . random_sample ( ) < self . prob_dict [ name ] : <EOL> data [ name ] [ : , i , ... ] = data [ name ] [ : , i , : , : : - <NUM_LIT:1> , : ] <EOL> yield data <EOL> class OneHot ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , vocab_size_dict ) : <EOL> """<STR_LIT>""" <EOL> DataIterator . __init__ ( self , iter . data_shapes , iter . length ) <EOL> for key in vocab_size_dict . keys ( ) : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> if not isinstance ( vocab_size_dict [ key ] , int ) : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> shape = iter . data_shapes [ key ] <EOL> if not ( shape [ - <NUM_LIT:1> ] == <NUM_LIT:1> and len ( shape ) == <NUM_LIT:3> ) : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> self . vocab_size_dict = vocab_size_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler = None ) : <EOL> for data in self . iter ( handler ) : <EOL> for name in self . vocab_size_dict . keys ( ) : <EOL> vocab_size = self . vocab_size_dict [ name ] <EOL> new_data = np . eye ( vocab_size , dtype = np . bool ) [ data [ name ] ] <EOL> new_data = new_data . reshape ( ( new_data . shape [ <NUM_LIT:0> ] , <EOL> new_data . shape [ <NUM_LIT:1> ] , <EOL> new_data . shape [ <NUM_LIT:3> ] ) ) <EOL> data [ name ] = new_data <EOL> yield data <EOL> class MultiHot ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , vocab_size_dict ) : <EOL> """<STR_LIT>""" <EOL> DataIterator . __init__ ( self , iter . data_shapes , iter . length ) <EOL> for key in vocab_size_dict . keys ( ) : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> if not isinstance ( vocab_size_dict [ key ] , int ) : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> shape = iter . data_shapes [ key ] <EOL> if not len ( shape ) == <NUM_LIT:3> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> self . vocab_size_dict = vocab_size_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler ) : <EOL> for data in self . iter ( handler ) : <EOL> for name in self . vocab_size_dict . keys ( ) : <EOL> vocab_size = self . vocab_size_dict [ name ] <EOL> new_data = np . eye ( vocab_size , dtype = np . bool ) [ data [ name ] ] . max ( <NUM_LIT:2> ) <EOL> data [ name ] = new_data <EOL> yield data <EOL> class Pad ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , size_dict , value_dict = None ) : <EOL> """<STR_LIT>""" <EOL> super ( Pad , self ) . __init__ ( iter . data_shapes , iter . length ) <EOL> if value_dict is not None : <EOL> if set ( size_dict . keys ( ) ) != set ( value_dict . keys ( ) ) : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> for key in size_dict . keys ( ) : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> if len ( iter . data_shapes [ key ] ) != <NUM_LIT:5> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> self . value_dict = { } if value_dict is None else value_dict <EOL> self . size_dict = size_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler = None ) : <EOL> for data in self . iter ( handler ) : <EOL> for name in self . size_dict . keys ( ) : <EOL> assert isinstance ( data [ name ] , np . ndarray ) <EOL> t , b , h , w , c = data [ name ] . shape <EOL> size = self . size_dict [ name ] <EOL> val = self . value_dict . get ( name , <NUM_LIT:0.0> ) <EOL> new_data = val * np . ones ( ( t , b , h + <NUM_LIT:2> * size , w + <NUM_LIT:2> * size , c ) ) <EOL> new_data [ : , : , size : - size , size : - size , : ] = data [ name ] <EOL> data [ name ] = new_data <EOL> yield data <EOL> class RandomCrop ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , iter , shape_dict ) : <EOL> """<STR_LIT>""" <EOL> super ( RandomCrop , self ) . __init__ ( iter . data_shapes , iter . length ) <EOL> for key , val in shape_dict . items ( ) : <EOL> if key not in iter . data_shapes : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT:}>" . format ( key , iter . data_shapes . keys ( ) ) ) <EOL> if not ( isinstance ( val , tuple ) and len ( val ) == <NUM_LIT:2> ) : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> data_shape = iter . data_shapes [ key ] <EOL> if len ( data_shape ) != <NUM_LIT:5> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> if val [ <NUM_LIT:0> ] > data_shape [ <NUM_LIT:2> ] or val [ <NUM_LIT:0> ] < <NUM_LIT:0> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> if val [ <NUM_LIT:1> ] > data_shape [ <NUM_LIT:3> ] or val [ <NUM_LIT:1> ] < <NUM_LIT:0> : <EOL> raise IteratorValidationError ( "<STR_LIT>" ) <EOL> self . shape_dict = shape_dict <EOL> self . iter = iter <EOL> def __call__ ( self , handler = None ) : <EOL> for data in self . iter ( handler ) : <EOL> for name in self . shape_dict . keys ( ) : <EOL> assert isinstance ( data [ name ] , np . ndarray ) <EOL> t , n , h , w , c = data [ name ] . shape <EOL> crop_h , crop_w = self . shape_dict [ name ] <EOL> max_r = h - crop_h <EOL> max_c = w - crop_w <EOL> row_indices = self . rnd . random_integers ( <NUM_LIT:0> , max_r , n ) <EOL> col_indices = self . rnd . random_integers ( <NUM_LIT:0> , max_c , n ) <EOL> cropped = np . zeros ( ( t , n , crop_h , crop_w , c ) ) <EOL> _crop_images ( data [ name ] , crop_h , crop_w , row_indices , <EOL> col_indices , cropped ) <EOL> data [ name ] = cropped <EOL> yield data <EOL> class Undivided ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , ** named_data ) : <EOL> """<STR_LIT>""" <EOL> _assert_correct_data_format ( named_data ) <EOL> data_shapes = { n : v . shape for n , v in named_data . items ( ) } <EOL> super ( Undivided , self ) . __init__ ( data_shapes , <NUM_LIT:1> ) <EOL> self . data = named_data <EOL> self . total_size = int ( sum ( d . size for d in self . data . values ( ) ) ) <EOL> def __call__ ( self , handler = None ) : <EOL> yield self . data <EOL> class Minibatches ( DataIterator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , batch_size = <NUM_LIT:1> , shuffle = True , cut_according_to = '<STR_LIT>' , <EOL> ** named_data ) : <EOL> """<STR_LIT>""" <EOL> nr_sequences , time_steps = _assert_correct_data_format ( named_data ) <EOL> data_shapes = { n : v . shape for n , v in named_data . items ( ) } <EOL> nr_batches = int ( math . ceil ( nr_sequences / batch_size ) ) <EOL> super ( Minibatches , self ) . __init__ ( data_shapes , nr_batches ) <EOL> self . data = named_data <EOL> self . shuffle = shuffle <EOL> self . batch_size = batch_size <EOL> if isinstance ( cut_according_to , ( six . string_types , type ( None ) ) ) : <EOL> if cut_according_to in named_data : <EOL> self . seq_lens = _calculate_lengths_from_mask ( <EOL> named_data [ cut_according_to ] ) <EOL> else : <EOL> self . seq_lens = time_steps * np . ones ( nr_sequences , <EOL> dtype = np . int ) <EOL> else : <EOL> self . seq_lens = np . array ( cut_according_to ) <EOL> assert self . seq_lens . shape == ( nr_sequences , ) <EOL> self . sample_size = int ( <EOL> sum ( d . shape [ <NUM_LIT:0> ] * np . prod ( d . shape [ <NUM_LIT:2> : ] ) * batch_size <EOL> for d in self . data . values ( ) ) ) <EOL> def __call__ ( self , handler = None ) : <EOL> indices = np . arange ( self . length ) <EOL> if self . shuffle : <EOL> self . rnd . shuffle ( indices ) <EOL> for idx in indices : <EOL> batch_slice = slice ( idx * self . batch_size , <EOL> ( idx + <NUM_LIT:1> ) * self . batch_size ) <EOL> time_slice = slice ( None , np . max ( self . seq_lens [ batch_slice ] ) ) <EOL> data = { k : v [ time_slice , batch_slice ] <EOL> for k , v in self . data . items ( ) } <EOL> yield data <EOL> def _assert_correct_data_format ( named_data ) : <EOL> nr_sequences = { } <EOL> nr_timesteps = { } <EOL> for name , data in named_data . items ( ) : <EOL> if not hasattr ( data , '<STR_LIT>' ) : <EOL> raise IteratorValidationError ( <EOL> "<STR_LIT>" . format ( name ) <EOL> ) <EOL> if len ( data . shape ) < <NUM_LIT:3> : <EOL> raise IteratorValidationError ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> nr_sequences [ name ] = data . shape [ <NUM_LIT:1> ] <EOL> nr_timesteps [ name ] = data . shape [ <NUM_LIT:0> ] <EOL> if min ( nr_sequences . values ( ) ) != max ( nr_sequences . values ( ) ) : <EOL> raise IteratorValidationError ( <EOL> '<STR_LIT>' <EOL> . format ( nr_sequences ) ) <EOL> if min ( nr_timesteps . values ( ) ) != max ( nr_timesteps . values ( ) ) : <EOL> raise IteratorValidationError ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' . format ( nr_timesteps ) ) <EOL> return int ( min ( nr_sequences . values ( ) ) ) , min ( nr_timesteps . values ( ) ) <EOL> def _calculate_lengths_from_mask ( mask ) : <EOL> assert mask . shape [ <NUM_LIT:2> : ] == ( <NUM_LIT:1> , ) <EOL> b = mask [ : , : , <NUM_LIT:0> ] != <NUM_LIT:0> <EOL> lengths = mask . shape [ <NUM_LIT:0> ] - b [ : : - <NUM_LIT:1> ] . argmax ( axis = <NUM_LIT:0> ) <EOL> lengths [ b . max ( axis = <NUM_LIT:0> ) == <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> return lengths </s>
