如何在TensorFlow中批量处理任意尺寸图像？求等比例缩放到短边500的方法

Hey there! Let's tackle your two questions step by step, and tie in the practical constraints from your research setup to make this actionable.

The core challenge here is that uneven image dimensions can't be directly stacked into a fixed-shape tensor. Here are two reliable approaches tailored to avoid the memory/speed issues you mentioned with oversized images:

Option 1: Use Dynamic Shapes with tf.data (For Fully Convolutional Models) #

If your model is fully convolutional (no fixed-size dense layers), you can work directly with variable-size images. Use tf.data to load and batch them without forcing a uniform size:

import tensorflow as tf

def preprocess_image(image_path):
    # Load and decode image
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img, channels=3)
    
    # Resize oversized images to short side = 500 (matches your Caffe workflow)
    h, w = tf.shape(img)[0], tf.shape(img)[1]
    short_side = tf.minimum(h, w)
    scale = 500.0 / tf.cast(short_side, tf.float32)
    new_h = tf.cast(tf.cast(h, tf.float32) * scale, tf.int32)
    new_w = tf.cast(tf.cast(w, tf.float32) * scale, tf.int32)
    img = tf.image.resize(img, (new_h, new_w), method=tf.image.ResizeMethod.BILINEAR)
    
    # Normalize (match ImageNet preprocessing for pretrained conv layers)
    img = tf.keras.applications.resnet50.preprocess_input(img)
    return img

# Build dataset
image_paths = ["path/to/img1.jpg", "path/to/img2.jpg", ...]
dataset = tf.data.Dataset.from_tensor_slices(image_paths)
dataset = dataset.map(preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)

# Batch directly (works for models that accept dynamic input shapes)
dataset = dataset.batch(batch_size=128)

Option 2: Padded Batches (For Models with Fixed-Size Layers) #

If your model has fixed-shape layers (e.g., dense heads), use padded_batch to pad images to the largest dimensions in each batch:

# After preprocessing, pad images to match batch's max size
dataset = dataset.padded_batch(
    batch_size=128,
    padding_values=0.0,  # Pad with black pixels (adjust based on your normalization)
    padded_shapes=tf.TensorShape([None, None, 3])
)

Bonus: Match Your Research Training Setup #

To replicate your Caffe training config in TensorFlow:

• Use SGD with momentum and learning rate scheduling:

  # Learning rate scheduler: decay every 20k iterations
  lr_scheduler = tf.keras.optimizers.schedules.StepDecay(
      initial_learning_rate=1e-3,
      step_size=20000,
      decay_rate=0.1  # Adjust decay factor to match your needs
  )
  optimizer = tf.keras.optimizers.SGD(
      learning_rate=lr_scheduler,
      momentum=0.9,
      weight_decay=5e-4
  )
  

• Load pretrained conv layers: Use tf.keras.applications (e.g., ResNet50(weights='imagenet', include_top=False)) and fine-tune on your AVA dataset.

Caffe's built-in data layers don't natively support aspect-ratio preserving resizing to a target short side, but you can implement this easily with a custom Python data layer:

Step 1: Write the Image Loading/Resizing Logic #

import cv2
import numpy as np
import caffe

class AspectResizeDataLayer(caffe.Layer):
    def setup(self, bottom, top):
        # Parse params from prototxt
        params = eval(self.param_str)
        self.target_short_side = params.get('target_short_side', 500)
        self.mean = np.array(params.get('mean', [103.939, 116.779, 123.68]))  # ImageNet mean
        self.batch_size = 128
        # Add your own logic to load image paths (e.g., from a text file)
        self.image_path_list = self.load_image_paths()
        self.current_idx = 0

    def get_next_image_path(self):
        path = self.image_path_list[self.current_idx]
        self.current_idx = (self.current_idx + 1) % len(self.image_path_list)
        return path

    def reshape(self, bottom, top):
        self.images = []
        max_h, max_w = 0, 0
        for _ in range(self.batch_size):
            img_path = self.get_next_image_path()
            img = cv2.imread(img_path)
            
            # Resize to short side = 500, preserve aspect ratio
            h, w = img.shape[:2]
            scale = self.target_short_side / min(h, w)
            new_h, new_w = int(h * scale), int(w * scale)
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
            
            # Track max dimensions for padding (if needed)
            max_h = max(max_h, new_h)
            max_w = max(max_w, new_w)
            
            # Convert to Caffe's format (BGR -> RGB, channel-first, subtract mean)
            img = img[:, :, ::-1].astype(np.float32)
            img -= self.mean
            img = img.transpose((2, 0, 1))
            self.images.append(img)
        
        # Reshape top blob to fit padded batch
        top[0].reshape(self.batch_size, 3, max_h, max_w)

    def forward(self, bottom, top):
        # Copy and pad images to top blob
        for i in range(self.batch_size):
            h, w = self.images[i].shape[1], self.images[i].shape[2]
            top[0].data[i, :, :h, :w] = self.images[i]

    def backward(self, top, propagate_down, bottom):
        # No backward pass needed for data layers
        pass

Step 2: Update Your Caffe Prototxt #

Add the custom layer to your train/test prototxt:

layer {
  name: "data"
  type: "Python"
  top: "data"
  top: "label"
  python_param {
    module: "your_module_name"
    layer: "AspectResizeDataLayer"
    param_str: '{"target_short_side": 500, "mean": [103.939, 116.779, 123.68]}'
  }
}

This approach exactly matches your research workflow: it keeps the original aspect ratio, caps the short side at 500 to avoid memory issues, and works seamlessly with your SGD optimizer and pretrained conv layers.

内容的提问来源于stack exchange，提问作者Ryan L.