import os, random, pickle
import numpy as np
import scipy.misc

def get_mask_path(path):
    name, ext = os.path.splitext(path)
    return name + "_mask.png"

max_image_size = 256

images_and_masks = []

for root, subdirs, files in os.walk('train_images'):
    for file in files:
        path = os.path.join(root, file)
        if path.endswith("_mask.png"): continue

        print("loading image %s"%path)
        
        image = scipy.misc.imread(path)
        mask = scipy.misc.imread(get_mask_path(path), flatten=True)

        image_height, image_width, image_channels = image.shape
        
        # discard alpha channel
        if image_channels == 4:
            image = image[:, :, :3]
        
        image_height, image_width, image_channels = image.shape
        assert(image_channels == 3)
        
        scale = max_image_size/float(max(image_width, image_height))

        image = scipy.misc.imresize(image, scale)
        mask = scipy.misc.imresize(mask, scale)

        height, width = mask.shape
        mask = mask.reshape((height, width, 1))

        image = image.astype(np.float32)/255.0
        mask = mask.astype(np.float32)/255.0

        images_and_masks.append((image, mask))

print("writing %d images"%len(images_and_masks))

with open("data.pickle", "wb") as f:
    pickle.dump(images_and_masks, f)
