Использование tensorflow.data для создания набора данных изображений и нескольких этикеток

Question

Я пытаюсь обучить нейронную сеть рисовать ограничивающую рамку вокруг объекта. Я сам сгенерировал данные, изображения 256x256 rgb и пять меток на изображение (два угла ограничивающей рамки + компонент вращения). Чтобы не исчерпать память при обучении сети с использованием python 3.7.6, tensorflow 2.0 и keras, я загружал за раз только небольшое количество изображений. Затем сеть обучилась на них, а затем загружается новый набор изображений. Все это происходит последовательно (и я не очень хороший программист, вероятно, не очень эффективно), что оставило меня с тем, что кажется довольно серьезным узким местом из-за способа загрузки изображений и этикеток. Имена изображений даны в виде числовых значений и в настоящее время сохраняются как.jpg, а мои ярлыки хранятся в текстовом файле, где каждая строка соответствует имени изображения.

Чтобы уменьшить / устранить узкое место, я прочитал о tf.data и попытался следовать примеру в https://www.tensorflow.org/tutorials/load_data/images#using_tfdata_for_finer_control. Однако в этих примерах рассматривается классификация, поэтому метки создаются по-другому. Я пытался изменить этот код вот так

import numpy as np
import tensorflow as tf
import os

height=256
width=256

AUTOTUNE = tf.data.experimental.AUTOTUNE
image_count=len(os.listdir('images_train'))

list_ds = tf.data.Dataset.list_files(str('images_train/*'), shuffle=False)
list_ds = list_ds.shuffle(image_count, reshuffle_each_iteration=False)

print('\n')
for f in list_ds.take(5):
    
    print(f.numpy())
   
def decode_img(img):
    # convert the compressed string to a 3D uint8 tensor
    img = tf.image.decode_jpeg(img, channels=3)
    # resize the image to the desired size
    return tf.image.resize(img, [height, width])  

#This is the function I cannot figure out how to write
def get_label(file_path):
    labels=np.loadtxt('labels_train.txt', delimiter=',')
    # convert the path to a list of path components
    parts = tf.strings.split(file_path, os.path.sep)
    
    #somehow I would like to extract the name of the image file and then take the numerical part and 
    #return the corresponding row
    return labels[0,:]
   
def process_path(file_path):
    label = get_label(file_path)
    # load the raw data from the file as a string
    img = tf.io.read_file(file_path)
    img = decode_img(img)
    return img, label

train_ds = list_ds.map(process_path, num_parallel_calls=AUTOTUNE)

Когда я просто возвращаю строку из файла, остальная часть скрипта, кажется, работает нормально, но я не могу понять, как сделать так, чтобы каждое изображение было связано с соответствующей меткой. Чтобы извлечь имя изображения в get_label() функция, которую я пытался использовать parts.numpy() что дает только это AttributeError: 'Tensor' object has no attribute 'numpy'.

Я пытался понять это в течение нескольких дней и не смог найти сообщение, которое точно описывает ту же проблему.

Как эффективно решить эту проблему, не будучи опытным программистом? Мы очень ценим все, что указывает мне правильное направление.