Python 实现 GPU 加速图像处理

1. 使用 PyTorch 实现 GPU 加速的卷积滤波(如边缘检测)

python 复制代码
import torch
import torch.nn as nn
import cv2
import numpy as np

# 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 读取图像并转换为 PyTorch 张量
image = cv2.imread("input.jpg")  # 读取 BGR 格式图像
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转为 RGB
image_tensor = torch.from_numpy(image).float().permute(2, 0, 1)  # HWC -> CHW
image_tensor = image_tensor.unsqueeze(0).to(device)  # 添加 batch 维度并移至 GPU

# 定义边缘检测卷积核(Sobel算子)
conv_layer = nn.Conv2d(
    in_channels=3,
    out_channels=3,
    kernel_size=3,
    bias=False,
    padding=1
).to(device)

# 手动设置 Sobel 核权重(示例,仅作用于水平边缘)
sobel_kernel = torch.tensor([
    [[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]],  # Red 通道
    [[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]],  # Green 通道
    [[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]],  # Blue 通道
], dtype=torch.float32).repeat(3, 1, 1, 1).to(device)

conv_layer.weight.data = sobel_kernel

# 执行卷积操作(GPU加速)
with torch.no_grad():
    output_tensor = conv_layer(image_tensor)

# 将结果转换回 numpy 并保存
output = output_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
output = np.clip(output, 0, 255).astype(np.uint8)
cv2.imwrite("edge_detection_gpu.jpg", cv2.cvtColor(output, cv2.COLOR_RGB2BGR))

2. 使用 OpenCV 的 CUDA 模块加速高斯模糊

python 复制代码
import cv2
import time

# 检查 OpenCV 是否支持 CUDA
print("CUDA devices:", cv2.cuda.getCudaEnabledDeviceCount())

# 读取图像并上传到 GPU
image = cv2.imread("input.jpg")
gpu_image = cv2.cuda_GpuMat()
gpu_image.upload(image)

# 创建 GPU 加速的高斯滤波器
gaussian_filter = cv2.cuda.createGaussianFilter(
    cv2.CV_8UC3,  # 输入类型 (8-bit unsigned, 3 channels)
    cv2.CV_8UC3,  # 输出类型
    (15, 15),      # 核大小
    0              # Sigma(自动计算)
)

# 执行滤波(重复多次测试速度)
start_time = time.time()
for _ in range(100):  # 重复 100 次模拟大数据量
    gpu_blur = gaussian_filter.apply(gpu_image)
end_time = time.time()

# 下载结果到 CPU 并保存
result = gpu_blur.download()
print(f"GPU Time: {end_time - start_time:.4f} seconds")
cv2.imwrite("blur_gpu.jpg", result)

3. 使用 CuPy 加速图像傅里叶变换

python 复制代码
import cupy as cp
import cv2
import numpy as np
import time

# 读取图像并转为灰度
image = cv2.imread("input.jpg", cv2.IMREAD_GRAYSCALE)

# 将 numpy 数组转为 CuPy 数组(上传到 GPU)
image_gpu = cp.asarray(image)

# 快速傅里叶变换(FFT)和逆变换(IFFT)
start_time = time.time()
fft_gpu = cp.fft.fft2(image_gpu)
fft_shift = cp.fft.fftshift(fft_gpu)
magnitude_spectrum = cp.log(cp.abs(fft_shift))
end_time = time.time()

# 将结果转回 CPU
magnitude_cpu = cp.asnumpy(magnitude_spectrum)
print(f"GPU FFT Time: {end_time - start_time:.4f} seconds")

# 归一化并保存频谱图
magnitude_cpu = cv2.normalize(magnitude_cpu, None, 0, 255, cv2.NORM_MINMAX)
cv2.imwrite("fft_spectrum_gpu.jpg", magnitude_cpu.astype(np.uint8))

4. 使用 Numba 编写自定义 GPU 核函数(图像反色)

python 复制代码
from numba import cuda
import numpy as np
import cv2
import time

# 读取图像
image = cv2.imread("input.jpg")
height, width, channels = image.shape

# 定义 GPU 核函数
@cuda.jit
def invert_colors_kernel(image):
    x, y = cuda.grid(2)
    if x < image.shape[0] and y < image.shape[1]:
        for c in range(3):  # 遍历 RGB 通道
            image[x, y, c] = 255 - image[x, y, c]

# 将图像上传到 GPU
image_gpu = cuda.to_device(image)

# 配置线程和块
threads_per_block = (16, 16)
blocks_per_grid_x = (height + threads_per_block[0] - 1) // threads_per_block[0]
blocks_per_grid_y = (width + threads_per_block[1] - 1) // threads_per_block[1]
blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y)

# 执行核函数
start_time = time.time()
invert_colors_kernel[blocks_per_grid, threads_per_block](image_gpu)
cuda.synchronize()  # 等待 GPU 完成
end_time = time.time()

# 下载结果并保存
image_cpu = image_gpu.copy_to_host()
print(f"GPU Invert Time: {end_time - start_time:.6f} seconds")
cv2.imwrite("inverted_gpu.jpg", image_cpu)

5. 使用 PyTorch 实现实时风格迁移(GPU加速)

python 复制代码
import torch
import torchvision.models as models
from torchvision import transforms
from PIL import Image

# 加载预训练模型到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.vgg19(pretrained=True).features.to(device).eval()

# 图像预处理
preprocess = transforms.Compose([
    transforms.Resize(512),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载内容图像和风格图像
content_image = Image.open("content.jpg")
style_image = Image.open("style.jpg")

# 将图像转为张量并移至 GPU
content_tensor = preprocess(content_image).unsqueeze(0).to(device)
style_tensor = preprocess(style_image).unsqueeze(0).to(device)

# 定义风格迁移函数(示例,需完整实现损失计算和优化)
def style_transfer(model, content_input, style_input, iterations=500):
    # 创建可优化图像
    input_image = content_input.clone().requires_grad_(True)
    
    # 定义优化器
    optimizer = torch.optim.LBFGS([input_image])
    
    # 风格迁移循环
    for i in range(iterations):
        def closure():
            optimizer.zero_grad()
            # 提取特征并计算损失(需实现具体细节)
            # ...
            return total_loss
        
        optimizer.step(closure)
    
    return input_image

# 执行风格迁移(需补充完整代码)
output_image = style_transfer(model, content_tensor, style_tensor)

# 后处理并保存结果
output_image = output_image.squeeze().cpu().detach()
output_image = transforms.ToPILImage()(output_image)
output_image.save("style_transfer_gpu.jpg")

关键说明

  1. 硬件依赖:需 NVIDIA GPU 并安装正确版本的 CUDA 和 cuDNN。

  2. 库安装

    bash 复制代码
    pip install torch torchvision opencv-python-headless cupy numba
  3. 性能对比:与 CPU 版本相比,GPU 加速通常快 10-100 倍(取决于任务复杂度)。

  4. 适用场景

    • PyTorch:适合深度学习相关的图像处理(如 GAN、超分辨率)。
    • OpenCV CUDA:适合传统图像处理加速(滤波、特征提取)。
    • CuPy/Numba:适合自定义数值计算或科研算法。
相关推荐
Yharim11 分钟前
wxPython官方文档中文翻译 - TipProvider概述
python
猿榜20 分钟前
深入浅出 Python 面向对象编程
javascript·python
jimin_callon20 分钟前
VBA第三十九期 VBA自动把工作表转换成PPT
python·excel·vba·wps·deepseek
ak啊23 分钟前
Python 实现对Excel 文件的统计与处理
python
Linhieng23 分钟前
AI使用经验分享:完成合并PDF功能
python
开开心心就好28 分钟前
手机不同App音量自动调节软件
网络·windows·python·安全·智能手机·电脑·音视频
depend29 分钟前
AI生成应用:图片批量重命名工具 - 自动化整理您的图片库
python
wumule30 分钟前
文本转语音TTS模型部署时UniDic(日语分词包)无法下载问题
python
齐尹秦31 分钟前
CSS 创建与使用学习笔记
人工智能·python·tensorflow
JavaEdge在掘金31 分钟前
深入理解 Java Final:从基础到线程安全与性能优化
python