如何使用训练好的SegFormer模型对单张图像进行预测？

你当前的代码存在几个关键问题，导致无法正确完成单张图像预测：

• 模型加载逻辑错误：直接用SegformerForSemanticSegmentation()加载PyTorch Lightning生成的checkpoint行不通，因为这个checkpoint保存的是你自定义的SegformerFinetuner（LightningModule）的完整状态，而非基础SegFormer模型的独立权重。
• 单张图像加载方式错误：ImageFolder是用于加载结构化分类数据集的工具（要求目录下分设子类文件夹），不适合加载单张独立图像。
• 缺少设备匹配处理：训练时使用了GPU，预测时需确保模型和图像处于同一计算设备上。

正确的预测步骤与代码 #

1. 加载训练好的Lightning模型

直接通过自定义的SegformerFinetuner从checkpoint加载完整模型：

import pytorch_lightning as pl
import torch

# 加载完整的LightningModule，需传入初始化时的必要参数（根据你的SegformerFinetuner类定义调整）
segformer_finetuner = SegformerFinetuner.load_from_checkpoint(
    checkpoint_path='lightning_logs/version_33/checkpoints/epoch=151-step=304.ckpt',
    id2label=train_dataset.id2label
)
segformer_finetuner.eval()
segformer_finetuner.freeze()  # 冻结参数，避免意外更新

# 匹配训练时的计算设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
segformer_finetuner.to(device)

2. 正确加载并预处理单张图像

使用PIL加载单张图像，应用与训练时完全一致的预处理流程：

from PIL import Image
import torchvision.transforms as transforms

# 注意：这里的transform要和训练数据集使用的完全一致，比如训练时做了resize就必须保留
transform = transforms.Compose([
    transforms.Resize((512, 512)),  # 替换为你训练时用的图像尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载单张图像并处理
img_path = 'path/to/your/image.jpg'
img = Image.open(img_path).convert('RGB')  # 确保图像为RGB格式
img_tensor = transform(img).unsqueeze(0).to(device)  # 添加batch维度并移到对应设备

3. 执行预测并处理结果

SegFormer输出的是logits，需要转换为类别掩码，并上采样至原图尺寸（如果模型输出尺寸小于原图）：

with torch.no_grad():
    output = segformer_finetuner(img_tensor)  # 直接调用LightningModule的forward方法预测

# 提取logits并上采样到原图尺寸
logits = output.logits  # 根据你的模型输出结构调整，比如output['logits']
upsampled_logits = torch.nn.functional.interpolate(
    logits,
    size=img.size[::-1],  # PIL图像size是(w, h)，PyTorch需要(h, w)，所以反转
    mode='bilinear',
    align_corners=False
)

# 生成像素级类别掩码
pred_mask = torch.argmax(upsampled_logits, dim=1).squeeze(0).cpu().numpy()

# 可选：将类别ID转换为标签
pred_labels = [segformer_finetuner.id2label[id] for id in pred_mask.flatten()]

内容的提问来源于stack exchange，提问作者Alex