-
Notifications
You must be signed in to change notification settings - Fork 0
/
eval_on_test.py
81 lines (63 loc) · 2.38 KB
/
eval_on_test.py
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
import argparse
import os
from test import engine
import torch
from unet import Unet
from torch.utils.data.dataloader import DataLoader
from eval import eval_net
from melanomia_dataset import MelanomiaDataset
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument('-mp',
'--model_path',
type=str,
required = True,
dest='model_path')
parser.add_argument('-i',
'--images',
type=str,
required = True,
dest='image_directory')
parser.add_argument('-m',
'--masks',
type=str,
dest='mask_directory',
required = True)
parser.add_argument('-o',
'--output',
type=str,
required=True,
dest='output_directory',
help='Output destination of segmentations')
parser.add_argument('-p',
'--padding',
type=bool,
default=True,
help='Add padding in the convolutions')
parser.add_argument('-s',
'--scale',
type=float,
default='.5',
dest='scale',
help='Scaling of images')
parser.add_argument('-u',
'--upscale',
type=bool,
default=True,
help='Upscaling of segmenatation to input resolution')
return parser.parse_args()
def run_predictions(args):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = MelanomiaDataset(args.image_directory,
args.mask_directory,
args.scale)
loader = DataLoader(dataset, num_workers=1, pin_memory=True)
net = Unet(addPadding=args.padding)
net.load_state_dict(
torch.load(args.model_path, map_location=device)
)
net.to(device=device)
print(engine(net, loader, device, args.output_directory, 'final_model'))
if __name__ == '__main__':
args = parse_args()
run_predictions(args)