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:适合自定义数值计算或科研算法。
相关推荐
数据智能老司机6 小时前
精通 Python 设计模式——分布式系统模式
python·设计模式·架构
数据智能老司机7 小时前
精通 Python 设计模式——并发与异步模式
python·设计模式·编程语言
数据智能老司机7 小时前
精通 Python 设计模式——测试模式
python·设计模式·架构
数据智能老司机7 小时前
精通 Python 设计模式——性能模式
python·设计模式·架构
c8i7 小时前
drf初步梳理
python·django
每日AI新事件7 小时前
python的异步函数
python
这里有鱼汤8 小时前
miniQMT下载历史行情数据太慢怎么办?一招提速10倍!
前端·python
databook17 小时前
Manim实现脉冲闪烁特效
后端·python·动效
程序设计实验室17 小时前
2025年了,在 Django 之外,Python Web 框架还能怎么选?
python
倔强青铜三19 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试