Image_Classification_PyTorch/test.py at master · FandosA/Image_Classification_PyTorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import os
import json
import torch
import numpy as np
import torch.nn as nn
from PIL import Image
import configargparse
from ConvNet import ConvNet
from torchvision.transforms import transforms


def predict(dataset_pred, img_name, transformer):

    image = Image.open(dataset_pred + '/' + img_name)
    image_tensor = torch.unsqueeze(transformer(image), dim=0)
    image_tensor = image_tensor.to(device)

    predictions = model(image_tensor)

    softmax = nn.Softmax(dim=1)
    predictions_softmax = softmax(predictions)
    predictions_softmax = torch.squeeze(predictions_softmax)
    predictions_softmax = predictions_softmax.cpu().detach().numpy()

    max_index = np.argmax(predictions_softmax)
    probability = predictions_softmax[max_index]
    predicted_class = classes[max_index]

    return predicted_class, probability, predictions_softmax


if __name__ == "__main__":

    # Select parameters for predictions
    p = configargparse.ArgumentParser()
    p.add_argument('--dataset_pred', type=str, default='seg_pred', help='Dataset predictions path.')
    p.add_argument('--log_dir', type=str, default='image_classification_ConvNet', help='Name of the folder to load the model.')
    p.add_argument('--checkpoint', type=str, default='checkpoint_18_best.pth',help='Checkpoint path')
    p.add_argument('--device', type=str, default='gpu', help='Choose the device: "gpu" or "cpu"')
    opt = p.parse_args()

    assert os.path.isdir(opt.log_dir), 'The folder log_dir does not exists'


    # Select device
    if opt.device == 'gpu' and torch.cuda.is_available():
        device = torch.device("cuda:0")
        print('Device assigned: GPU (' + torch.cuda.get_device_name(device) + ')\n')
    else:
        device = torch.device("cpu")
        if not torch.cuda.is_available() and opt.device == 'gpu':
            print('GPU not available, device assigned: CPU\n')
        else:
            print('Device assigned: CPU\n')


    classes = os.listdir("seg_test")
    images = os.listdir(opt.dataset_pred)
    num_images = len(images)

    transformer = transforms.Compose([
        transforms.Resize((150,150)),
        transforms.ToTensor(),  # 0-255 to 0-1, numpy to tensors
        # formula: output[channel] = (input[channel] - mean[channel]) / std[channel]
        transforms.Normalize([0.5,0.5,0.5], # 0-1 to [-1,1]
                            [0.5,0.5,0.5])
    ])

    model = ConvNet(num_classes=len(classes)).to(device)
    state_dict = torch.load(opt.log_dir + "/checkpoints/" + opt.checkpoint, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()

    predictions_dict={
        # Name of the image
        "image": [],

        # Class predicted for the image
        "class": [],

        # Probability that the image is the predicted class
        "probability": [],

        # Probabilities predicted for each class
        "probabilities": []
    }

    for i, img_name in enumerate(images):

        print('\rPredicting images... [' + str(i+1) + '/' + str(num_images) + ']', end='')

        predicted_class, probability, predictions = predict(opt.dataset_pred, img_name, transformer)

        predictions_dict["image"].append(img_name)
        predictions_dict["class"].append(predicted_class)
        predictions_dict["probability"].append(str(probability))
        predictions_dict["probabilities"].append(predictions.tolist())


    with open(opt.log_dir + '/predictions.json', "w") as fp:
        json.dump(predictions_dict, fp, indent=4)