def get_dataset(configs):
"""
This function get raw dataset
"""
from utils.datasets.z_dataset import z
# todo: add transform
train_set = z('train', configs)
val_set = z('val', configs)
test_set = z('test', configs, tta=True, tta_size=10)
return train_set, val_set, test_set
def main():
with open("./configs/z_config.json") as f:
configs = json.load(f)
acc = 0.0
state = torch.load("./saved/checkpoints/{}".format(checkpoint_name))
from models import resmasking_dropout1
model = resmasking_dropout1
model = model(in_channels=3, num_classes=7).cuda()
model.load_state_dict(state["net"])
model.eval()
correct = 0
total = 0
all_target = []
all_output = []
test_set = z("test", configs, tta=True, tta_size=8)
with torch.no_grad():
for idx in tqdm(range(len(test_set)), total=len(test_set),
leave=False):
images, targets = test_set[idx]
images = make_batch(images)
images = images.cuda(non_blocking=True)
outputs = model(images).cpu()
outputs = F.softmax(outputs, 1)
# outputs.shape [tta_size, 7]
outputs = torch.sum(outputs, 0)
outputs = torch.argmax(outputs, 0)
outputs = outputs.item()
# targets = targets.item()
total += 1
correct += outputs == targets
all_target.append(targets)
all_output.append(outputs)
# acc = 100. * correct / total
# print("Accuracy {:.03f}".format(acc))
all_target = np.array(all_target)
all_output = np.array(all_output)
matrix = confusion_matrix(all_target, all_output)
np.set_printoptions(precision=2)
# plt.figure(figsize=(5, 5))
plot_confusion_matrix(
matrix,
classes=class_names,
normalize=True,
# title='{} \n Accuracc: {:.03f}'.format(checkpoint_name, acc)
title="Residual Masking Network",
)
# plt.show()
# plt.savefig('cm_{}.png'.format(checkpoint_name))
plt.savefig("./saved/cm/cm_vemo_{}.pdf".format(checkpoint_name))
plt.close()
print("plt saved at ./saved/cm/cm_vemo_{}.pdf".format(checkpoint_name))