如何使用 PyTorch CNN 保存图像路径答案

【问题标题】：How to save image paths using PyTorch CNN如何使用 PyTorch CNN 保存图像路径
【发布时间】：2020-03-14 11:56:02
【问题描述】：

我一直在 PyTorch 中使用 CNN，我需要为每个类保存图像路径及其相关的预测概率（在这种情况下，类是通过或失败）。这是我将 preds 保存到数据框的代码：

preds_df = pd.DataFrame()   
class_labels = []


model_ft.eval()
for i, (inputs, labels) in enumerate(dataloaders['train']):
    inputs = inputs.to(device)
    labels = labels.to(device)
    class_labels.append(labels.tolist())
    output = model_ft(inputs)

    sm = torch.nn.Softmax()
    probabilities = sm(output) 
    arr = probabilities.data.cpu().numpy()
    df = pd.DataFrame(arr)

    preds_df = preds_df.append(df)


preds_df['prediction'] = preds_df.idxmax(axis=1)
class_list = [item for sublist in class_labels for item in sublist]
preds_df['label'] = class_list
preds_df.columns = ['pass (0)', 'fail (1)', 'prediction', 'label']


preds_df.to_csv('./zoom17CNN_preds.csv')

如何在数据加载器中保存每个文件的图像路径？谢谢！

【问题讨论】：

见this 和this。基本上，您必须手动编写 Dataset 以便它显式返回文件名以及您的图像和标签。

标签： python conv-neural-network pytorch

【解决方案1】：

感谢@akshayk07！我最终遍历了我的图像目录并以这种方式保存了当前图像名称和 preds：

directory = "./get_preds/fail"

for filename in os.listdir(directory):
    school_id = filename[0:6]
    ids.append(school_id)
    to_open = "./get_preds/fail/" + filename
    png = Image.open(to_open)

    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),         
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    img_t = transform(png)
    batch_t = torch.unsqueeze(img_t, 0)
    model_ft.eval()
    out = model_ft(batch_t)

    _, index = torch.max(out, 1)
    percentage = torch.nn.functional.softmax(out, dim=1)[0]
    class0.append(percentage[0].tolist())
    class1.append(percentage[1].tolist())

df['school_id'] = ids
df['pass (0)'] = class0
df['fail (1)'] = class1
df['label'] = 1

【讨论】：