for file_name in oslistdirdata_dir # 加载图像文件 image_path = ospathjoindata_dir file_name image_pathsappendimage_path image = Imageopenimage_pathconvertRGB

要保存每张图片前五个标签的概率，您需要进行以下修改：

1. 创建一个空列表top5_predictions来保存前五个标签的概率。
2. 在预测过程中，将输出的概率值也保存到top5_predictions列表中。
3. 修改写入文件的逻辑，将每张图片的路径、真实标签和前五个标签的概率一起保存到文件中。

下面是修改后的代码：

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

data_dir = 'your_data_directory'  # 替换为您的数据目录

# 加载模型和预处理变换
model = your_model  # 替换为您的模型
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # 标准化
])

image_paths = []
predictions = []
top5_predictions = []  # 保存前五个标签的概率

for file_name in os.listdir(data_dir):
    # 加载图像文件
    image_path = os.path.join(data_dir, file_name)
    image_paths.append(image_path)
    image = Image.open(image_path).convert('RGB')

    # 对图像进行预处理
    image_tensor = transform(image)
    image_batch = image_tensor.unsqueeze(0)

    # 在模型上进行预测
    with torch.no_grad():
        image_batch = image_batch.cuda(non_blocking=True)
        output = model(image_batch)
        _, predicted = torch.max(output, 1)
        predictions.extend(predicted.cpu().numpy())
        top5_probs, top5_indices = torch.topk(torch.softmax(output, dim=1), k=5)
        top5_predictions.append(top5_probs.cpu().numpy())

with open('real_labels.txt', 'r') as real_file:
    real_labels = [line.strip() for line in real_file.readlines()]

with open('predictions.txt', 'w') as file:
    for i in range(len(predictions)):
        file.write(f"{image_paths[i]} {real_labels[predictions[i]]}\n")
        file.write(f"Top 5 Predictions: {[(real_labels[top5_indices[i][j]], top5_probs[i][j]) for j in range(5)]}\n")

这样，您就可以保存每张图片的前五个标签的概率了。请记得将your_data_directory和your_model替换为您自己的数据目录和模型