64 lines
1.6 KiB
Python
64 lines
1.6 KiB
Python
import torch
|
|
from torchvision import datasets
|
|
from torch.utils.data import DataLoader
|
|
from sklearn.metrics import confusion_matrix
|
|
import numpy as np
|
|
|
|
from preprocess import TRANSFORM
|
|
from settings import (
|
|
RAW_DATASET_PATH,
|
|
MODEL_FILENAME,
|
|
get_balancing_sampler,
|
|
SetNameTp,
|
|
OurCNN,
|
|
)
|
|
|
|
from sys import argv
|
|
|
|
TARGET: SetNameTp = "test"
|
|
|
|
if __name__ == "__main__":
|
|
|
|
target: SetNameTp = TARGET
|
|
|
|
if len(argv) > 1:
|
|
idx = SetNameTp.__args__.index(argv[1])
|
|
target = SetNameTp.__args__[idx]
|
|
|
|
print(f"Testing on {target} set")
|
|
|
|
model = OurCNN()
|
|
model.load_state_dict(torch.load(MODEL_FILENAME))
|
|
|
|
test_dataset = datasets.ImageFolder(
|
|
root=RAW_DATASET_PATH + target, transform=TRANSFORM
|
|
)
|
|
|
|
test_loader = DataLoader(
|
|
test_dataset,
|
|
batch_size=128,
|
|
num_workers=4,
|
|
sampler=get_balancing_sampler(test_dataset),
|
|
)
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
predicted_y = np.array([])
|
|
true_y = np.array([])
|
|
|
|
for images, labels in test_loader:
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
with torch.no_grad():
|
|
outputs = model(images)
|
|
_, predicted = torch.max(outputs, 1)
|
|
predicted_y = np.append(predicted_y, predicted.cpu().numpy())
|
|
true_y = np.append(true_y, labels.cpu().numpy())
|
|
|
|
accuracy = (predicted_y == true_y).sum().item() / predicted_y.size
|
|
print(f"Accuracy: {accuracy:.2f}")
|
|
|
|
print(confusion_matrix(true_y, predicted_y))
|