from PIL import Image
import torchvision.transforms as T
from torch.autograd import Variable as V
import torch as t
model = model.cuda()#导入网络模型
model.eval()
model.load_state_dict(t.load('./models/new/model.dat'))#加载训练好的模型文件
import numpy as np
def result_(res):
if res==0:
return 'airplane'
elif res==1:
return 'ship'
elif res==2:
return 'bridge'
elif res==3:
return 'oilcan'
elif res==4:
return 'build'
else:
return 'Nan'
trans=T.Compose([
T.Scale(325),
T.CenterCrop(299),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
trans_gray=T.Compose([
T.Scale(325),
T.CenterCrop(299),
T.ToTensor(),
T.Normalize((0.1307,), (0.3081,))
])
import csv
import cv2
#stag_01_submit
# csvFile = open("test_submit.csv", "w")
csvFile = open("stag_01_submit.csv", "w") #创建csv文件
writer = csv.writer(csvFile) #创建写的对象
#先写入columns_name
writer.writerow(["id","label","ship","bridge","airplane","build","oilcan"]) #写入列的名称
#读入图片
# test_root = './test/'
test_root = './stag_01/'
img_test=os.listdir(test_root)
for i in range(len(img_test)):
rd_img = cv2.imread(test_root+img_test[i])
img = Image.open(test_root+img_test[i])
img = img.convert('RGB')
print(test_root+img_test[i])
print(rd_img.shape)
input=trans(img)
input=input.unsqueeze(0)#这里经过转换后输出的input格式是[C,H,W],网络输入还需要增加一维批量大小B
#增加一维，输出的img格式为[1,C,H,W]
input = V(input.cuda())
score = model(input)#将图片输入网络得到输出
probability = t.nn.functional.softmax(score,dim=1)#计算softmax，即该图片属于各类的概率
max_value,index = t.max(probability,1)#找到最大概率对应的索引号，该图片即为该索引号对应的类别
class_index = result_(index)
# print(class_index)
probability=np.round(probability.cpu().detach().numpy(),3)
# print(probability[0][0])
writer.writerow([img_test[i],class_index,probability[0][1],probability[0][2],probability[0][0],probability[0][4],probability[0][3]])
csvFile.close()
import pandas as pd
data = pd.read_csv('test_submit.csv')
