<s> import os <EOL> import glob <EOL> import sys <EOL> import yaml <EOL> import scipy . misc <EOL> import numpy as np <EOL> import hickle as hkl <EOL> def get_img ( img_name , img_size = <NUM_LIT> , batch_size = <NUM_LIT> ) : <EOL> target_shape = ( img_size , img_size , <NUM_LIT:3> ) <EOL> img = scipy . misc . imread ( img_name ) <EOL> assert img . dtype == '<STR_LIT>' , img_name <EOL> if len ( img . shape ) == <NUM_LIT:2> : <EOL> img = scipy . misc . imresize ( img , ( img_size , img_size ) ) <EOL> img = np . asarray ( [ img , img , img ] ) <EOL> else : <EOL> if img . shape [ <NUM_LIT:2> ] > <NUM_LIT:3> : <EOL> img = img [ : , : , : <NUM_LIT:3> ] <EOL> img = scipy . misc . imresize ( img , target_shape ) <EOL> img = np . rollaxis ( img , <NUM_LIT:2> ) <EOL> if img . shape [ <NUM_LIT:0> ] != <NUM_LIT:3> : <EOL> print img_name <EOL> return img <EOL> def save_batches ( file_list , tar_dir , img_size = <NUM_LIT> , batch_size = <NUM_LIT> , <EOL> flag_avg = False , num_sub_batch = <NUM_LIT:1> ) : <EOL> '''<STR_LIT>''' <EOL> if not os . path . exists ( tar_dir ) : <EOL> os . makedirs ( tar_dir ) <EOL> img_batch = np . zeros ( ( <NUM_LIT:3> , img_size , img_size , batch_size ) , np . uint8 ) <EOL> if flag_avg : <EOL> img_sum = np . zeros ( ( <NUM_LIT:3> , img_size , img_size ) ) <EOL> batch_count = <NUM_LIT:0> <EOL> count = <NUM_LIT:0> <EOL> for file_name in file_list : <EOL> img_batch [ : , : , : , count % batch_size ] = get_img ( file_name , img_size = img_size , batch_size = batch_size ) <EOL> count += <NUM_LIT:1> <EOL> if count % batch_size == <NUM_LIT:0> : <EOL> batch_count += <NUM_LIT:1> <EOL> if flag_avg : <EOL> img_sum += img_batch . mean ( axis = <NUM_LIT:3> ) <EOL> if num_sub_batch == <NUM_LIT:1> : <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch , os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> elif num_sub_batch == <NUM_LIT:2> : <EOL> half_size = batch_size / <NUM_LIT:2> <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch [ : , : , : , : half_size ] , <EOL> os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch [ : , : , : , half_size : ] , <EOL> os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> elif num_sub_batch == <NUM_LIT:4> : <EOL> q1 = batch_size / <NUM_LIT:4> <EOL> q2 = batch_size / <NUM_LIT:2> <EOL> q3 = batch_size / <NUM_LIT:4> * <NUM_LIT:3> <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch [ : , : , : , : q1 ] , <EOL> os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch [ : , : , : , q1 : q2 ] , <EOL> os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch [ : , : , : , q2 : q3 ] , <EOL> os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> save_name = '<STR_LIT>' % ( batch_count - <NUM_LIT:1> ) + '<STR_LIT>' <EOL> hkl . dump ( img_batch [ : , : , : , q3 : ] , <EOL> os . path . join ( tar_dir , save_name ) , mode = '<STR_LIT:w>' ) <EOL> else : <EOL> NotImplementedError ( "<STR_LIT>" ) <EOL> return img_sum / batch_count if flag_avg else None <EOL> def get_train_filenames ( src_train_dir , misc_dir , seed = <NUM_LIT:1> ) : <EOL> if os . path . exists ( os . path . join ( misc_dir , '<STR_LIT>' ) ) : <EOL> return np . load ( os . path . join ( misc_dir , '<STR_LIT>' ) ) <EOL> if not os . path . exists ( misc_dir ) : <EOL> os . makedirs ( misc_dir ) <EOL> print '<STR_LIT>' <EOL> subfolders = [ name for name in os . listdir ( src_train_dir ) <EOL> if os . path . isdir ( os . path . join ( src_train_dir , name ) ) ] <EOL> train_filenames = [ ] <EOL> for subfolder in subfolders : <EOL> train_filenames += glob . glob ( src_train_dir + subfolder + '<STR_LIT>' ) <EOL> train_filenames = np . asarray ( sorted ( train_filenames ) ) <EOL> np . random . seed ( seed ) <EOL> np . random . shuffle ( train_filenames ) <EOL> np . save ( os . path . join ( misc_dir , '<STR_LIT>' ) , <EOL> train_filenames ) <EOL> return train_filenames <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> with open ( '<STR_LIT>' , '<STR_LIT:r>' ) as f : <EOL> paths = yaml . load ( f ) <EOL> train_img_dir = paths [ '<STR_LIT>' ] <EOL> val_img_dir = paths [ '<STR_LIT>' ] <EOL> misc_dir = paths [ '<STR_LIT>' ] <EOL> if len ( sys . argv ) < <NUM_LIT:2> : <EOL> gen_type = '<STR_LIT>' <EOL> else : <EOL> gen_type = sys . argv [ <NUM_LIT:1> ] <EOL> if gen_type == '<STR_LIT>' : <EOL> print '<STR_LIT>' <EOL> elif gen_type == '<STR_LIT>' : <EOL> print '<STR_LIT>' <EOL> else : <EOL> NotImplementedError ( "<STR_LIT>" ) <EOL> train_filenames = get_train_filenames ( train_img_dir , misc_dir ) <EOL> val_filenames = np . asarray ( sorted ( glob . glob ( val_img_dir + '<STR_LIT>' ) ) ) <EOL> img_size = <NUM_LIT> <EOL> batch_size = <NUM_LIT> <EOL> if gen_type == '<STR_LIT>' : <EOL> train_filenames = train_filenames [ : <NUM_LIT> ] <EOL> val_filenames = val_filenames [ : <NUM_LIT> ] <EOL> for num_sub_batch in [ <NUM_LIT:1> , <NUM_LIT:2> ] : <EOL> tar_train_dir = paths [ '<STR_LIT>' ] <EOL> tar_val_dir = paths [ '<STR_LIT>' ] <EOL> tar_train_dir += '<STR_LIT>' + str ( batch_size ) + '<STR_LIT>' + str ( batch_size / num_sub_batch ) <EOL> tar_val_dir += '<STR_LIT>' + str ( batch_size ) + '<STR_LIT>' + str ( batch_size / num_sub_batch ) <EOL> img_mean = save_batches ( train_filenames , tar_train_dir , <EOL> img_size = img_size , batch_size = batch_size , <EOL> flag_avg = True , num_sub_batch = num_sub_batch ) <EOL> np . save ( os . path . join ( misc_dir , '<STR_LIT>' ) , img_mean ) <EOL> save_batches ( val_filenames , tar_val_dir , <EOL> img_size = img_size , batch_size = batch_size , <EOL> num_sub_batch = num_sub_batch ) </s>
