U-NET模型测试代码

以下是一个使用 U-Net 模型进行测试的示例代码，使用 PyTorch 框架：

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

class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        
        # Encoder
        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        
        # Decoder
        self.upconv1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.conv5 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.conv6 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
        self.upconv3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.conv7 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
        self.conv8 = nn.Conv2d(64, out_channels, kernel_size=1)
        
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        # Encoder
        x1 = self.relu(self.conv1(x))
        x2 = self.relu(self.conv2(x1))
        x3 = self.relu(self.conv3(x2))
        x4 = self.relu(self.conv4(x3))
        
        # Decoder
        x = self.relu(self.upconv1(x4))
        x = torch.cat([x, x3], dim=1)
        x = self.relu(self.conv5(x))
        x = self.relu(self.upconv2(x))
        x = torch.cat([x, x2], dim=1)
        x = self.relu(self.conv6(x))
        x = self.relu(self.upconv3(x))
        x = torch.cat([x, x1], dim=1)
        x = self.relu(self.conv7(x))
        x = self.conv8(x)
        
        return x

# 测试模型
model = UNet(3, 1) # 输入通道数为3，输出通道数为1
model.load_state_dict(torch.load('model.pth')) # 加载保存的模型参数
model.eval() # 切换到评估模式

image = Image.open('test.jpg') # 读取测试图像
image = image.convert('RGB') # 转换为RGB格式
image = image.resize((256, 256)) # 缩放为256x256大小
image = torch.tensor([transforms.ToTensor()(image)]) # 转换为张量

with torch.no_grad():
    output = model(image) # 模型输出
    output = output.sigmoid() # 将输出映射到0~1之间
    output = (output > 0.5).float() # 将输出二值化
    
output = output.squeeze().cpu().numpy() # 转换为NumPy数组并去除批次维度
output = (output * 255).astype('uint8') # 将输出转换为0~255之间的整数

Image.fromarray(output).save('output.jpg') # 保存输出图像

在这个示例代码中，我们首先定义了一个 U-Net 模型，并加载了保存的模型参数。然后，我们读取一张测试图像，并将其转换为 PyTorch 张量。我们使用 `with torch.no_grad()` 上下文管理器来关闭自动求导，以减少内存占用和计算时间。

我们将测试图像输入到模型中，得到模型的输出。由于模型输出的值域为 $(-\infty, \infty)$，我们使用 `sigmoid()` 函数将其映射到 $(0, 1)$，然后使用阈值为 $0.5$ 的二值化操作将其转换为二值图像。最后，我们将输出转换为 NumPy 数组，并将其保存为图像文件。

在实际应用中，我们可以使用测试代码来评估模型的性能，并调整模型的超参数以获得更好的结果。