Оптимизируйте использование памяти с помощью plt.imread ()

Question

Я работаю с большим количеством изображений (120 КБ), каждое изображение является компонентом RGB + желтый одного изображения (поэтому у меня есть 30 000 уникальных изображений, все разбиты на 4 изображения: одно для Красного, Зеленого, Синего иЖелтый)

Для каждого идентификатора изображения я объединяю 4 компонента (RGB + желтый) в массив (M, N, 4) (где M и N - размеры изображения).

Я работаю со следующим кодом:

import pandas as pd
import numpy as np
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
from os import listdir

train_labels_data = pd.read_csv('/Documents/train.csv')

def merge_rgb(img_id, colours=['red','blue','green','yellow'], path = 'train'):

"""
For each images, returns an array of shape (M,N,4) 
where each dimension in the 4 are red, blue, green and yellow.
"""
    merged_colour_img = []

    for colour in colours:
        full_path = path + '/' + img_id + '_' + colour + '.png'
        colour_img=mpimg.imread(full_path)
        merged_colour_img.append(colour_img)    

    merged_colour_img = np.dstack((merged_colour_img))
    return merged_colour_img


def train_data_label(train_labels_data):

"""
From the train_labels csv file, create a list of labels, and create a large 
array for the train data in same order.
"""
    train_ids = [img_id for img_id in train_labels_data['Id']]
    train_labels = [label for label in train_labels_data['Target']]

    print ('Labels and Ids collected')

    train_data = []

    i=0
    for img_id in train_ids:

        print ('Merging Image')
        train_data_img = merge_rgb (img_id)
        print ('Merging done, appending the (M,N,4) array to a list')
        train_data.append(train_data_img)
        i += 1
        print ('Done appending, going to next image')
        print(i)

    print('Stacking all images in one big array')
    train_data = np.stack(train_data)

    return train_labels, train_data


train_labels, train_data = train_data_label(train_labels_data)


# SAVE OUTPUT

data_pickle_train = pickle.dumps(train_data)
data = open("/Documents/train_data.pkl","wb")
data.write(data_pickle_train)
data.close()

data_pickle_train_labels = pickle.dumps(train_labels)
data = open("/Documents/train_data_labels.pkl","wb")
data.write(data_pickle_train_labels)
data.close()

Однако этот код использует много памяти и вылетает на полпути до обработки всех изображений.Поскольку я работаю с изображениями, я подозреваю, что мог бы улучшить функцию merge_rgb, любой совет, как это сделать?

Спасибо,