1 import os, json, argparse
  2 from threading import Thread
  3 from Queue import Queue
  4 
  5 import numpy as np
  6 from scipy.misc import imread, imresize
  7 import h5py
  8 
  9 """
 10 Create an HDF5 file of images for training a feedforward style transfer model.
 11 """
 12 
 13 parser = argparse.ArgumentParser()
 14 parser.add_argument('--train_dir', default='/media/wangxiao/WangXiao_Dataset/CoCo/train2014')
 15 parser.add_argument('--val_dir', default='/media/wangxiao/WangXiao_Dataset/CoCo/val2014')
 16 parser.add_argument('--output_file', default='/media/wangxiao/WangXiao_Dataset/CoCo/coco-256.h5')
 17 parser.add_argument('--height', type=int, default=256)
 18 parser.add_argument('--width', type=int, default=256)
 19 parser.add_argument('--max_images', type=int, default=-1)
 20 parser.add_argument('--num_workers', type=int, default=2)
 21 parser.add_argument('--include_val', type=int, default=1)
 22 parser.add_argument('--max_resize', default=16, type=int)
 23 args = parser.parse_args()
 24 
 25 
 26 def add_data(h5_file, image_dir, prefix, args):
 27   # Make a list of all images in the source directory
 28   image_list = []
 29   image_extensions = {'.jpg', '.jpeg', '.JPG', '.JPEG', '.png', '.PNG'}
 30   for filename in os.listdir(image_dir):
 31     ext = os.path.splitext(filename)[1]
 32     if ext in image_extensions:
 33       image_list.append(os.path.join(image_dir, filename))
 34   num_images = len(image_list)
 35 
 36   # Resize all images and copy them into the hdf5 file
 37   # We'll bravely try multithreading
 38   dset_name = os.path.join(prefix, 'images')
 39   dset_size = (num_images, 3, args.height, args.width)
 40   imgs_dset = h5_file.create_dataset(dset_name, dset_size, np.uint8)
 41   
 42   # input_queue stores (idx, filename) tuples,
 43   # output_queue stores (idx, resized_img) tuples
 44   input_queue = Queue()
 45   output_queue = Queue()
 46   
 47   # Read workers pull images off disk and resize them
 48   def read_worker():
 49     while True:
 50       idx, filename = input_queue.get()
 51       img = imread(filename)
 52       try:
 53         # First crop the image so its size is a multiple of max_resize
 54         H, W = img.shape[0], img.shape[1]
 55         H_crop = H - H % args.max_resize
 56         W_crop = W - W % args.max_resize
 57         img = img[:H_crop, :W_crop]
 58         img = imresize(img, (args.height, args.width))
 59       except (ValueError, IndexError) as e:
 60         print filename
 61         print img.shape, img.dtype
 62         print e
 63       input_queue.task_done()
 64       output_queue.put((idx, img))
 65   
 66   # Write workers write resized images to the hdf5 file
 67   def write_worker():
 68     num_written = 0
 69     while True:
 70       idx, img = output_queue.get()
 71       if img.ndim == 3:
 72         # RGB image, transpose from H x W x C to C x H x W
 73         imgs_dset[idx] = img.transpose(2, 0, 1)
 74       elif img.ndim == 2:
 75         # Grayscale image; it is H x W so broadcasting to C x H x W will just copy
 76         # grayscale values into all channels.
 77         imgs_dset[idx] = img
 78       output_queue.task_done()
 79       num_written = num_written + 1
 80       if num_written % 100 == 0:
 81         print 'Copied %d / %d images' % (num_written, num_images)
 82   
 83   # Start the read workers.
 84   for i in xrange(args.num_workers):
 85     t = Thread(target=read_worker)
 86     t.daemon = True
 87     t.start()
 88     
 89   # h5py locks internally, so we can only use a single write worker =(
 90   t = Thread(target=write_worker)
 91   t.daemon = True
 92   t.start()
 93     
 94   for idx, filename in enumerate(image_list):
 95     if args.max_images > 0 and idx >= args.max_images: break
 96     input_queue.put((idx, filename))
 97     
 98   input_queue.join()
 99   output_queue.join()
100   
101   
102   
103 if __name__ == '__main__':
104   
105   with h5py.File(args.output_file, 'w') as f:
106     add_data(f, args.train_dir, 'train2014', args)
107 
108     if args.include_val != 0:
109       add_data(f, args.val_dir, 'val2014', args)