from PIL import Image
import os
import pickle
from glob import glob
import numpy as np


"""
This code can pickle all images of the folder into one pickle file
"""

# set up directory for pickling and pickle name

SOURCE_DIRECTORY_L = "out_F/"
PICKLE_IMAGES_L = "labels_ext.p"

path_list_L = glob(os.path.join(SOURCE_DIRECTORY_L, "*.png"))
i = 0
list1 = []
with open(PICKLE_IMAGES_L, "wb") as f:
    for file_name in path_list_L:
            image_L = Image.open(file_name)
            im_l = image_L.resize((160, 80))
            im_l = np.asarray(im_l)
            # print(im_l.shape)
            green = im_l[:, :, 1]
            # print(green.shape)
            green = green[..., np.newaxis]
            # print(green.shape)
            list1.insert(i, green)
            i = i+1
    print(len(list1))
    pickle.dump(list1, f)

# set up directory for pickling and pickle name

SOURCE_DIRECTORY = "img_F/"
PICKLE_IMAGES = "train_ext.p"

path_list = glob(os.path.join(SOURCE_DIRECTORY, "*.png"))
i = 0
list2 = []
with open(PICKLE_IMAGES, "wb") as f:
    for file_name in path_list:
        image_T = Image.open(file_name)
        im_t = image_T.resize((160, 80))
        im_t = np.asarray(im_t)
        list2.insert(i, im_t)
        i = i + 1
    print(len(list2))
    pickle.dump(list2, f)

