45 lines
1.9 KiB
Python
45 lines
1.9 KiB
Python
# 使用CUDA进行加速，但实际效果是减速 ：（
|
||
|
||
import os
|
||
import torch
|
||
from torchvision.io.image import read_image, write_png
|
||
from torchvision.transforms.functional import to_pil_image
|
||
from PIL import Image
|
||
|
||
# 设置允许的最大像素数，例如设置为1亿像素
|
||
Image.MAX_IMAGE_PIXELS = 1000000000
|
||
def convert_jpg_to_webp_cuda(input_folder, output_folder):
|
||
# 确保输出文件夹存在
|
||
os.makedirs(output_folder, exist_ok=True)
|
||
|
||
# 遍历输入文件夹中的所有文件
|
||
for filename in os.listdir(input_folder):
|
||
# 检查文件是否为jpg或JPEG格式
|
||
if filename.endswith(".jpg") or filename.endswith(".jpeg"):
|
||
# 构建完整的文件路径
|
||
input_path = os.path.join(input_folder, filename)
|
||
# 获取新的文件名，将.jpg/.jpeg替换为.webp
|
||
webp_filename = filename.rsplit('.', 1)[0] + '.webp'
|
||
output_path = os.path.join(output_folder, webp_filename)
|
||
|
||
try:
|
||
# 使用PyTorch读取图像，尝试使用CUDA
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
image_tensor = read_image(input_path).to(device)
|
||
|
||
# 注意：这里我们先转换为PIL Image，因为torchvision目前不直接支持WEBP输出
|
||
# 实际上，这一步会将处理移回CPU，但整个流程仍可能受益于CUDA加速的预处理
|
||
pil_image = to_pil_image(image_tensor.cpu())
|
||
|
||
# 使用Pillow保存为WEBP
|
||
pil_image.save(output_path, format="WEBP")
|
||
print(f"Converted {filename} to {webp_filename} using CUDA")
|
||
except Exception as e:
|
||
print(f"Cannot convert {filename}: {e}")
|
||
|
||
# 指定输入和输出文件夹
|
||
input_folder = 'input'
|
||
output_folder = 'output_cuda'
|
||
|
||
# 调用函数进行转换
|
||
convert_jpg_to_webp_cuda(input_folder, output_folder) |