def predict_picture(pic_path,model_path,dataset,device):
if dataset == 'char74k':
size = 20
num_classes = 26
else:
size = 28
num_classes = 10
transform = transforms.Compose([
transforms.Scale(size),
transforms.CenterCrop((size,size)),
transforms.ToTensor(),
transforms.Normalize((0.1307,),(0.3081,))
])
img = Image.open(pic_path)
img_tensor = transform(img)
img_tensor = img_tensor.unsqueeze(0)
img_tensor = img_tensor.to(device)
model = CapsNet(num_classes=num_classes,conv_in=3)
model = model.to(device)
model.eval()
state_dict = torch.load(model_path)
model.load_state_dict(state_dict)
score = model(img_tensor)
probs = nn.functional.softmax(score,dim=-1)
max_value,index = torch.max(probs,dim=-1)
print('The picture {} is:{}'.format(pic_path.split('/')[-1],char74k_id2labels[int(index)]))
def eval(model_path, result_path, dataProcessor, num_classes, epoch, device,
eval_type, logger):
capsNet = CapsNet(num_classes=num_classes, conv_in=3)
state_dict = torch.load(model_path + '/model_state_dict.pkl')
capsNet.load_state_dict(state_dict)
capsNet.to(device)
valid_scores, valid_predicts, valid_labels = predict(
capsNet, dataProcessor, num_classes, epoch, device, 'valid', logger)
test_scores, test_predicts, test_labels = predict(capsNet, dataProcessor,
num_classes, epoch,
device, 'test', logger)
store_results(test_predicts, test_labels, result_path)
with open(result_path + '/result.csv', 'w') as file:
file.write('valid scores,test scores\n')
file.write(str(valid_scores) + ',' + str(test_scores))
return valid_scores, test_scores