65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
import csv
|
|
import torch
|
|
import mlflow
|
|
|
|
from torchvision import transforms, datasets
|
|
from torch.utils.data import DataLoader
|
|
from train import Model
|
|
|
|
|
|
def get_data(IMG_SIZE: int, BATCH_SIZE: int):
|
|
testTransformer = transforms.Compose([
|
|
transforms.Resize(size = (IMG_SIZE, IMG_SIZE), antialias = True),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
testSet = datasets.ImageFolder(root = "./test", transform = testTransformer)
|
|
testLoader = DataLoader(testSet, batch_size = BATCH_SIZE, shuffle = False)
|
|
|
|
return testLoader
|
|
|
|
|
|
if __name__ == '__main__':
|
|
IMG_SIZE = 224
|
|
BATCH_SIZE = 32
|
|
MODEL_PATH = 'model.pth'
|
|
|
|
mlflow.set_experiment('Skin Cancer Detection - Test set')
|
|
mlflow.start_run()
|
|
|
|
names = {0: 'Benign', 1: 'Malignant'}
|
|
predictions = []
|
|
|
|
test_loader = get_data(IMG_SIZE, BATCH_SIZE)
|
|
labels = [names[i] for i in test_loader.dataset.targets]
|
|
|
|
model = Model()
|
|
model.load_state_dict(torch.load(MODEL_PATH))
|
|
|
|
test_correct, test_total = 0, 0
|
|
|
|
with torch.no_grad():
|
|
for i, data in enumerate(test_loader):
|
|
input, label = data
|
|
output = model(input)
|
|
_, predicted = torch.max(output.data, 1)
|
|
test_total += label.size(0)
|
|
test_correct += (predicted == label).sum().item()
|
|
|
|
predictions.extend(predicted.tolist())
|
|
|
|
test_accuracy = test_correct / test_total
|
|
|
|
mlflow.log_metric('test_accuracy', test_accuracy)
|
|
|
|
predictions = [names[pred] for pred in predictions]
|
|
|
|
print(f'Accuracy test {test_accuracy:.2%}')
|
|
|
|
with open('predictions.csv', 'w', newline='') as csvfile:
|
|
writer = csv.writer(csvfile)
|
|
writer.writerow(["Predictions", "Labels"])
|
|
|
|
for pred, label in zip(predictions, labels):
|
|
writer.writerow([pred, label]) |