41 lines
1.2 KiB
Python
Executable file
41 lines
1.2 KiB
Python
Executable file
import os
|
|
import sys
|
|
from PIL import Image
|
|
from StringIO import StringIO
|
|
from util import *
|
|
|
|
src = '/ais/gobi3/u/ilya/flickr_85/'
|
|
dst = '/ais/gobi3/u/kriz/flickr-85-1024/'
|
|
|
|
BATCH_SIZE = 2048
|
|
|
|
def save_batch(c_strings, c_sizes, c_labels, out_b):
|
|
pickle(os.path.join(dst, 'data_batch_%d' % out_b), (c_strings, c_sizes, c_labels))
|
|
|
|
return out_b + 1
|
|
if __name__ == "__main__":
|
|
c_strings = []
|
|
c_sizes = []
|
|
c_labels = []
|
|
out_b = 1
|
|
for b in xrange(977):
|
|
failed = 0
|
|
strings, sizes, labels = unpickle(os.path.join(src, '%s' % b))
|
|
for s,z,l in zip(strings, sizes, labels):
|
|
try:
|
|
im = Image.open(StringIO(s)).convert('RGB')
|
|
c_strings += [s]
|
|
c_sizes += [z]
|
|
c_labels += [l]
|
|
|
|
if len(c_strings) == BATCH_SIZE:
|
|
out_b = save_batch(c_strings, c_sizes, c_labels, out_b)
|
|
c_strings = []
|
|
c_sizes = []
|
|
c_labels = []
|
|
except IOError,e:
|
|
failed += 1
|
|
print "Batch %d failed: %d" % (b, failed)
|
|
|
|
if len(c_strings) > 0:
|
|
save_batch(c_strings, c_sizes, c_labels, out_b)
|