# 使用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)