pytorch黑白像素转换的代码

以下是一个使用 PyTorch 实现黑白图像转换的简单示例代码：

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image

# 定义黑白图像转换模型
class BWConverter(nn.Module):
    def __init__(self):
        super(BWConverter, self).__init__()
        self.conv = nn.Conv2d(1, 3, kernel_size=3, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

# 加载模型
model = BWConverter()

# 加载预训练模型参数
model.load_state_dict(torch.load('bw_converter.pth'))

# 加载图片并进行预处理
img = Image.open('example.jpg').convert('L')
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(256),
    transforms.ToTensor(),
])
img = transform(img)
img = img.unsqueeze(0)

# 执行黑白图像转换
output = model(img)
output = output.detach().squeeze().permute(1, 2, 0)

# 保存转换后的图片
output_img = transforms.ToPILImage()(output)
output_img.save('output.jpg')

在上述示例代码中，我们定义了一个名为 `BWConverter` 的简单神经网络模型，用于将黑白图像转换为彩色图像。该模型包含一个 `Conv2d` 层和一个 `ReLU` 激活函数层。我们使用预训练的模型参数来实现黑白图像转换，并使用 PyTorch 提供的图片预处理工具，将处理后的图片输入到模型中进行转换。最终，我们将转换后的图片保存到本地文件中。在运行代码之前，请确保已经安装了 PyTorch 和 PIL 库，并将图片文件和预训练模型文件放在正确的位置。