With the code below, I first construct a dataset by reading and processing some images and their respective segmentation maps, and then map them to perform some augmentations. In the final mapping stage, I create a class weight to feed the model fit. As a result of this code, the weight_map seems to iterate independently from the augmentation step. However, the weights must be produced from the augmented segmentation maps. How can I solve this issue?
tf_dataset = tf.data.Dataset.from_tensor_slices((train_im_files, train_seg_files))
tf_dataset = tf_dataset.map(process_path, num_parallel_calls=AUTOTUNE, deterministic=True)
tf_augmented = tf_dataset.map(augment_batch, num_parallel_calls=AUTOTUNE, deterministic=True)
weight_map = tf_augmented.map(make_weight_map_batch, num_parallel_calls=AUTOTUNE, deterministic=True)
net.model.fit(tf_augmented.batch(5), epochs=100, class_weight=make_weight_map_batch)