【Python】Pytorchで推論をして，ファイル名，Label名，Predictionの結果をcsv出力する．

PytorchでTraining済みのmodelを用いて推論を行う。
結果をtest dataのファイル名と一緒にcsv出力する。

まずは、下準備

# module import
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision

from torch.utils.data import Dataset, DataLoader,TensorDataset,random_split,SubsetRandomSampler, ConcatDataset
from torch.nn import functional as F

from torchvision import datasets,transforms
import torchvision.transforms as transforms

from sklearn.metrics import classification_report


# device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

ws = "ワークスペースフォルダを指定"
os.chdir(ws)

# test data フォルダ
test_data_dir = 'test image を入れているフォルダ'

Data loaderの準備

# data_transforms
data_transforms = {
        torchvision.transforms.Compose([
        #torchvision.transforms.RandomResizedCrop(224),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

test_dataset = torchvision.datasets.ImageFolder(test_data_dir,  transform=data_transforms)


# なぜか、これを経ないと走らない。。。たぶん違う書き方がある
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

test_dataset.transform=transform
#%%
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16)

modelの準備

model_path = "modelのpath"
model = torch.load(model_path)
model.eval()
model.to("device")   #cpuでの高速で推論は可能

ファイル名を抽出する。

# data_transforms
data_transforms = {
        torchvision.transforms.Compose([
        #torchvision.transforms.RandomResizedCrop(224),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

test_dataset = torchvision.datasets.ImageFolder(test_data_dir,  transform=data_transforms)


# なぜか、これを経ないと走らない。。。たぶん違う書き方がある
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

test_dataset.transform=transform
#%%
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16)

modelの準備

fn= []
# get file name　into list
test_data_fns = test_loader.dataset.samples
fn += [l[0] for l in test_data_fns]

推論の実行

pred = []
Y = []

for i, (x,y) in enumerate(test_loader):
    with torch.no_grad():
        output = model(x)
    pred += [int(l.argmax()) for l in output]
    Y += [int(l) for l in y]
    
print(classification_report(Y, pred))

dataframeにて、concatで結合してcsv出力

import pandas as pd

#dataframe
df_fn    = pd.Series(fn)
df_pred  = pd.Series(pred)
df_label = pd.Series(Y)


#dataframeを縦に結合
tmp_df = pd.concat([df_fn,df_pred,df_label],axis=1)
#dataframeをcolumnの名前を編集
tmp_df = tmp_df.rename(columns={0:'File name',1:'Prediction',2:'Label',3:'IC'})
#csvを出力
tmp_df.to_csv('Result_summary.csv', header=True, index=True)