Входной конвейер на основе tf.dataset с предшествующей маской для обработки видео в реальном времени

Question

Я хочу перейти на входной конвейер на основе tf.dataset. Все учебные пособия показывают, как выполнять классификацию или сегментацию изображений, и в этой задаче нет никаких проб, но я работаю над обработкой видео в реальном времени и мне нужно пройти предварительную маску с данными кадра, чтобы обеспечить временную согласованность при прогнозировании. В настоящее время у меня есть генератор на основе numpy для кера, в котором обучаемая модель используется для генерации маски (на основе предыдущего кадра) и используется в качестве предыдущей маски для текущего кадра (4 канала подаются в сеть). Предварительное создание маски происходит во время создания новой партии. Этот код работает хорошо. Я попробовал много подходов, как сделать то же самое, но с tf.dataset без успеха. Может быть, кто-то сталкивался с подобной проблемой и может поделиться решением?

Заранее спасибо!

Какой-то экспериментальный код

# Dummy network    
g = tf.Graph()
with g.as_default():
    inputs = tf.keras.Input(shape=(28,28,4)) 
    x = tf.keras.layers.Conv2D(filters=32, 
                               kernel_size=(3, 3), 
                               activation=tf.nn.relu)(inputs)
    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(x)
    x = tf.keras.layers.Conv2D(filters=64, 
                               kernel_size=(3, 3), 
                               activation=tf.nn.relu)(x)
    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(x)
    out = tf.keras.layers.Conv2D(filters=2, 
                               kernel_size=(3, 3), 
                               activation=tf.nn.relu)(x)
    model = tf.keras.Model(inputs=inputs, outputs=out)
    optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
    model.compile(loss='categorical_crossentropy',
                  optimizer=optimizer,
                  metrics=['accuracy'])
    model.summary()    

# Class for wrapping model 
class prior_mask_creator():

    def __init__(self, model_prior):
        self.model_prior = model_prior

    def make_prior_mask(self, x, y):
        h = tf.get_session_handle(x)
        h = K.get_session().run(h)
        v_1 = K.get_session().run([self.model_prior.layers[-1].output], feed_dict={inputs:h })
        return tf.cast(x, tf.float32), tf.cast(y, tf.float32)

# Image generator 
c_g = custom_generator(x_data, y_data, batch)

# Create class-wrapper
mask_creator = prior_mask_creator(model)

# Create dataset
ds_test = tf.data.Dataset.from_generator(c_g.gen,
                                         output_types=(tf.float32, tf.float32),
                                         output_shapes=(tf.TensorShape([None, None, 4]), tf.TensorShape([None, None, 2])))

ds_test = ds_test.map(random_crop)
# Add prior prediction
ds_test = ds_test.map(mask_creator.make_prior_mask)

ds_test = ds_test.batch(batch)

iterator = ds_test.make_initializable_iterator()
next_batch = iterator.get_next()

with g.as_default():
    with tf.Session(graph=g) as sess:
        K.set_session(sess)

        sess.run(tf.global_variables_initializer())
        sess.run(iterator.initializer)

        for i in range(epoch):
            v = sess.run(next_batch)