如何将蒙版应用于tf.data管道中的图像？

import tensorflow as tf
import numpy as np

original_image = tf.io.read_file("path/to/original_image.jpg")
original_image = tf.image.decode_jpeg(original_image, channels=3)
original_image = tf.image.convert_image_dtype(original_image, tf.float32)

mask_image = tf.io.read_file("path/to/mask_image.jpg")
mask_image = tf.image.decode_jpeg(mask_image, channels=1)
mask_image = tf.image.convert_image_dtype(mask_image, tf.float32)

masked_image = tf.multiply(original_image, mask_image)

def apply_mask(original_image_path, mask_image_path):
    original_image = tf.io.read_file(original_image_path)
    original_image = tf.image.decode_jpeg(original_image, channels=3)
    original_image = tf.image.convert_image_dtype(original_image, tf.float32)

    mask_image = tf.io.read_file(mask_image_path)
    mask_image = tf.image.decode_jpeg(mask_image, channels=1)
    mask_image = tf.image.convert_image_dtype(mask_image, tf.float32)

    masked_image = tf.multiply(original_image, mask_image)
    
    return masked_image

dataset = tf.data.Dataset.from_tensor_slices((original_image_paths, mask_image_paths))
dataset = dataset.map(lambda x, y: apply_mask(x, y))