TensorFlow动态拆分宽幅图像为方形块及数据集处理求助

Let's walk through fixing and expanding your code to meet your requirements: overlapping data reading/training, and splitting wide images into square blocks. Here's a complete, tested implementation with explanations:

Step 1: Imports & Core Functions #

First, we'll define helper functions to load images and split them into square blocks:

import tensorflow as tf

def load_and_preprocess_image(filename):
    """Read, decode, and normalize an image from disk."""
    # Read the raw image file
    img_raw = tf.io.read_file(filename)
    # Decode image (handles PNG, JPG, and other common formats)
    # expand_animations=False ensures we don't load GIFs as sequences
    img = tf.image.decode_image(img_raw, channels=3, expand_animations=False)
    # Convert to float32 and normalize to [0, 1] for stable training
    img = tf.cast(img, tf.float32) / 255.0
    return img

def split_into_square_blocks(image):
    """Split a wide image into square blocks (block size = image height)."""
    # Get image dimensions (height, width, channels)
    height = tf.shape(image)[0]
    width = tf.shape(image)[1]
    channels = tf.shape(image)[2]
    
    # Square block size matches the image's height
    block_size = height
    
    # Calculate how many blocks fit (round up to avoid cutting off content)
    num_blocks = tf.cast(tf.math.ceil(tf.cast(width, tf.float32) / block_size), tf.int32)
    
    # Pad the image width to make it evenly divisible by block size
    padded_width = num_blocks * block_size
    padded_image = tf.image.pad_to_bounding_box(image, 0, 0, height, padded_width)
    
    # Split the padded image into individual square blocks
    blocks = tf.reshape(padded_image, (num_blocks, height, block_size, channels))
    
    # Return blocks as a dataset so we can flatten them later
    return tf.data.Dataset.from_tensor_slices(blocks)

Step 2: Build the Efficient Data Pipeline #

Now we'll construct the dataset pipeline to overlap reading and training, and handle variable block counts:

def main(unused_argv):
    # Replace with your actual list of image filenames
    filenames = tf.constant(['im_01.png', 'im_02.png', 'im_03.png', 'im_04.png'])
    
    # 1. Create a dataset from filenames
    dataset = tf.data.Dataset.from_tensor_slices(filenames)
    
    # 2. Load images in parallel (speeds up data ingestion)
    dataset = dataset.map(
        load_and_preprocess_image,
        num_parallel_calls=tf.data.AUTOTUNE  # Let TensorFlow optimize parallelism
    )
    
    # 3. Split each image into blocks, then flatten into a dataset of blocks
    # flat_map is critical here: it converts each image's blocks into individual dataset elements
    dataset = dataset.flat_map(split_into_square_blocks)
    
    # 4. Optional: Resize blocks to a fixed size (required if your model expects uniform inputs)
    # If your images have varying heights, uncomment this to standardize block size
    # def resize_block(block):
    #     return tf.image.resize(block, (256, 256))  # Adjust to your model's input size
    # dataset = dataset.map(resize_block, num_parallel_calls=tf.data.AUTOTUNE)
    
    # 5. Shuffle, batch, and prefetch for training efficiency
    dataset = dataset.shuffle(buffer_size=100)  # Shuffle blocks to avoid order bias
    dataset = dataset.batch(32)  # Adjust batch size based on your GPU memory
    # Prefetch batches to overlap data processing with model training
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    
    # Test the pipeline (remove this in production)
    for batch_idx, batch in enumerate(dataset):
        print(f"Batch {batch_idx} shape: {batch.shape}")
        # Expected shape: (batch_size, block_height, block_height, 3)

if __name__ == '__main__':
    tf.compat.v1.app.run(main=main)

Key Features Explained #

• Overlapping Read/Training: The prefetch(tf.data.AUTOTUNE) call tells TensorFlow to prepare the next batch of data while the model is training on the current batch, eliminating idle time.
• Variable Block Counts: Using flat_map instead of map ensures that each square block becomes a separate element in the dataset, even if different images produce different numbers of blocks.
• Parallel Processing: num_parallel_calls=tf.data.AUTOTUNE optimizes the number of parallel threads for loading and processing images, maximizing CPU utilization.
• Padding for Full Coverage: We pad images to ensure no content is cut off when splitting into blocks. If you prefer to truncate instead, replace the padding step with tf.image.crop_to_bounding_box.

Troubleshooting Tips #

• Fixed Input Size: If your model requires uniform input dimensions, uncomment the resize step and set the target size to match your model's input layer.
• Memory Issues: If you're handling extremely large images, consider adding a tf.image.resize step before splitting to reduce the image size first, or use smaller batch sizes.
• Image Format Errors: Ensure all your images are in a format TensorFlow can decode (PNG, JPG, etc.). If you have mixed formats, add error handling in load_and_preprocess_image.

内容的提问来源于stack exchange，提问作者Harald