TensorFlow 2.0中实现图像与面部关键点标签同步随机旋转

当然有完美的解决方案！你的问题核心在于在tf.data.Dataset管道中混用了PIL（Python原生库）和TensorFlow操作，这在TF2的图模式下很容易出问题。我们可以用纯TensorFlow原生操作来实现随机旋转，同时保证图像和关键点同步变换，完全避开tf.contrib和PIL的兼容性问题。

1. 核心思路 #

• 用TensorFlow的tf.raw_ops.ImageProjectiveTransformV3实现任意角度图像旋转（原生TF操作，完美适配tf.data）；
• 推导旋转的坐标变换公式，手动计算关键点的旋转后坐标，保证和图像旋转完全同步。

2. 代码实现 #

导入依赖

import tensorflow as tf
import numpy as np

定义旋转变换工具函数

def get_rotation_transform(angle, image_size):
    """计算图像旋转的投影变换矩阵（适配TF的ProjectiveTransform格式）"""
    h, w = image_size
    theta = tf.cast(angle, tf.float32) * np.pi / 180.0  # 角度转弧度
    
    # 旋转核心参数
    cos_theta = tf.cos(theta)
    sin_theta = tf.sin(theta)
    
    # 计算平移量，保证旋转后图像中心不变
    tx = (w - w*cos_theta + h*sin_theta) / 2.0
    ty = (h - h*cos_theta - w*sin_theta) / 2.0
    
    # 构造TF要求的变换矩阵格式：[a0, a1, a2, b0, b1, b2, c0, c1]
    transform = tf.stack([
        cos_theta, -sin_theta, tx,
        sin_theta, cos_theta, ty,
        0.0, 0.0
    ])
    return transform

同步旋转图像与关键点

def rotate_image_and_keypoints(image, keypoints, angle, input_size):
    """同时旋转图像和对应的面部关键点"""
    # 图像预处理：解码、 resize、转灰度
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [input_size, input_size])
    image = tf.image.rgb_to_grayscale(image)
    
    # 旋转图像
    transform = get_rotation_transform(angle, [input_size, input_size])
    rotated_image = tf.raw_ops.ImageProjectiveTransformV3(
        images=tf.expand_dims(image, 0),  # 添加batch维度
        transforms=tf.expand_dims(transform, 0),
        output_size=[input_size, input_size],
        fill_value=0  # 旋转边缘填充0，可根据需求改为图像均值等
    )
    rotated_image = tf.squeeze(rotated_image, 0)  # 移除batch维度
    
    # 旋转关键点
    theta = tf.cast(angle, tf.float32) * np.pi / 180.0
    cos_theta = tf.cos(theta)
    sin_theta = tf.sin(theta)
    
    # 转换坐标到以图像中心为原点
    x_center = input_size / 2.0
    y_center = input_size / 2.0
    x_rel = keypoints[:, 0] - x_center
    y_rel = keypoints[:, 1] - y_center
    
    # 应用旋转公式
    x_rot = x_rel * cos_theta - y_rel * sin_theta
    y_rot = x_rel * sin_theta + y_rel * cos_theta
    
    # 转换回原坐标系
    rotated_keypoints = tf.stack([x_rot + x_center, y_rot + y_center], axis=1)
    
    return rotated_image, rotated_keypoints

构建tf.data管道

def load_and_preprocess_data(image_path, keypoints_path):
    # 读取图像文件
    image = tf.io.read_file(image_path)
    
    # 读取关键点（假设关键点存储为每行x,y的txt文件，可根据你的数据集格式调整）
    keypoints_raw = tf.io.read_file(keypoints_path)
    keypoints_raw = tf.strings.split(keypoints_raw, '\n')[:-1]  # 移除空行
    keypoints = tf.strings.split(keypoints_raw, ',')
    keypoints = tf.strings.to_number(keypoints, tf.float32)
    keypoints = tf.reshape(keypoints, [-1, 2])  # 转为[关键点数量, 2]的格式
    
    # 生成-60~60度的随机旋转角度
    angle = tf.random.uniform([], minval=-60, maxval=60, seed=0)
    
    # 同步旋转图像和关键点
    rotated_image, rotated_keypoints = rotate_image_and_keypoints(
        image, keypoints, angle, input_size=224  # 替换为你的输入图像尺寸
    )
    
    # 图像归一化（可选，根据你的模型需求调整）
    rotated_image = tf.cast(rotated_image, tf.float32) / 255.0
    
    return rotated_image, rotated_keypoints

# 构建数据集（替换为你的图像路径和关键点路径列表）
image_paths = ['face1.jpg', 'face2.jpg', ...]
keypoints_paths = ['face1_keypoints.txt', 'face2_keypoints.txt', ...]

dataset = tf.data.Dataset.from_tensor_slices((image_paths, keypoints_paths))
dataset = dataset.map(load_and_preprocess_data, num_parallel_calls=tf.data.AUTOTUNE)
# 后续可添加shuffle、batch、prefetch等优化操作
dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

3. 关键注意事项 #

• 关键点坐标格式：确保你的关键点是(x, y)对应图像的(宽度, 高度)，如果是(y, x)格式，需要调整旋转公式中的x/y顺序；
• 边缘填充：fill_value可根据需求修改，比如用图像的均值填充会比纯黑更自然；
• 性能优化：num_parallel_calls=tf.data.AUTOTUNE和prefetch能充分利用CPU资源，提升数据加载速度。

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