#!/usr/bin/env python3
import os
import numpy as np
import imageio
import zipfile

output_dir='/home/belarm/laotzu/source/nns/datasets/nist-sd19/parsed'
# output_dir='/home/ubuntu/efs/nist-sd19/parsed'

out_zips = {}

for i in range(0x30, 0x7b):
    out_zips["{:x}".format(i)] = zipfile.ZipFile(os.path.join(output_dir, '{:x}.zip'.format(i)), 'w')

zfile = zipfile.ZipFile('/home/belarm/Downloads/by_class.zip')
# zfile = zipfile.ZipFile('/home/ubuntu/by_class.zip')
old_img_class = ''

for f in zfile.namelist():
    if f.endswith('.png') and 'train' not in f:
        _, img_class, _, key  = f.split('/')
        img = imageio.imread(zfile.open(f).read())
        img_normalized = img.astype(np.bool)
        out_zips[img_class].writestr(key + '_raw', img.tobytes())
        out_zips[img_class].writestr(key + '_normalized', img_normalized.tobytes())
        if old_img_class != img_class:
            print(img_class)
            old_img_class = img_class
for key, val in out_zips:
    val.close()
zfile.close()