利用Python 脚本生成 .h5 文件
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)
