Как получить правильные размеры формы, используя альбументацию с tenorflow

Question

Я кодирую сверточный автоматический кодировщик.

Я использую альбументации для увеличения данных в моем коде тензорного потока, но у меня возникла грубая проблема с размерами формы.

Я получил форму (32, 1, 256, 256, 3) вместо (32, 256, 256, 3) на выходе конвейера train_dataset.

Вот мой код:

      def process_path(image_input, image_output):
  # load the raw data from the file as a string
    img_input = tf.io.read_file(image_input)
    img_input = tf.image.decode_jpeg(img_input, channels=3)

    img_output = tf.io.read_file(image_output)
    img_output = tf.image.decode_jpeg(img_output, channels=3)

    return img_input, img_output

def aug_fn(image):
    data = {"image": image}
    aug_data = transforms(**data)
    aug_img = aug_data["image"]
    aug_img = tf.cast(aug_img/255.0, tf.float32)
    aug_img = tf.image.resize(aug_img, size=[256, 256])
    return aug_img

def process_aug(img_input, img_output):
    aug_img_input = tf.numpy_function(func=aug_fn, inp=[img_input], Tout=[tf.float32])
    aug_img_output = tf.numpy_function(func=aug_fn, inp=[img_output], Tout=[tf.float32])
    return aug_img_input, aug_img_output

def set_shapes(img_input, img_output):
    img_input.set_shape((256, 256, 3))
    img_output.set_shape((256, 256, 3))
    return img_input, img_output

transforms = OneOf([CLAHE(clip_limit=2), IAASharpen(), IAAEmboss(), RandomBrightnessContrast()], p=0.3)
train_dataset = (tf.data.Dataset.from_tensor_slices((DIR_TRAIN + filenames_train, 
                                                         DIR_TRAIN + filenames_train))
                        .map(process_path, num_parallel_calls=AUTO)
                        .map(process_aug, num_parallel_calls=AUTO)
                        .map(set_shapes, num_parallel_calls=AUTO)
                        .batch(parameters['BATCH_SIZE'])
                        .prefetch(AUTO)
                    )
next(iter(train_dataset))

Я думаю, что эта проблема возникает из-за функции process_aug с комментарием функции карты в конвейере набора данных, но я не нашел, где именно проблема в функции process_aug.

Answer 1

Я отредактировал функцию, добавив tf.reshape(), но я надеюсь, что один из вас сможет объяснить, как избежать добавления функции tf.numpy_function в процесс измерения?

      def process_aug(img_input, img_output):
aug_img_input = tf.numpy_function(func=aug_fn, inp=[img_input, 256], Tout=[tf.float32])
    aug_img_input = tf.reshape(aug_img_input, [256, 256, 3])
    aug_img_output = tf.numpy_function(func=aug_fn, inp=[img_output, 256], Tout=[tf.float32])
    aug_img_output = tf.reshape(aug_img_output, [256, 256, 3])
    return aug_img_input, aug_img_output