将仿射变换应用于TensorFlow2中的图像

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

def apply_affine_transform(image, transform_matrix):
    image_shape = tf.shape(image)
    batch_size = image_shape[0]
    num_channels = image_shape[3]

    # Reshape the image to a 4D tensor
    image = tf.reshape(image, [batch_size, image_shape[1], image_shape[2], num_channels])

    # Apply the affine transformation
    transformed_image = tf.contrib.image.transform(image, transform_matrix)

    return transformed_image

# Load an example image
image_path = 'example.jpg'
image = tf.io.read_file(image_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.convert_image_dtype(image, tf.float32)

# Define the transformation matrix
transform_matrix = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)

# Apply the affine transformation
transformed_image = apply_affine_transform(image, transform_matrix)

# Display the original and transformed images
plt.subplot(1, 2, 1)
plt.imshow(image)
plt.title('Original Image')

plt.subplot(1, 2, 2)
plt.imshow(transformed_image)
plt.title('Transformed Image')

plt.show()