-
Notifications
You must be signed in to change notification settings - Fork 59
/
api.py
113 lines (92 loc) · 3.95 KB
/
api.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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import argparse
import os
import torch
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
from PIL import Image
import numpy as np
from bts.model import DynamicUNet
from bts.classifier import BrainTumorClassifier
def get_arguments():
"""Returns the command line arguments as a dict"""
parser = argparse.ArgumentParser()
parser.add_argument('--file', required=False, type=str,
help='Single input file name.')
parser.add_argument('--dir', required=False,
type=str, help='Directory name with input images')
parser.add_argument('--ofp', required=False,
type=str, help='Single output file path with name.Use this if using "file" flag.')
parser.add_argument('--odp', required=False,
type=str, help='Directory path for output images.Use this if using "dir" flag.')
args = parser.parse_args()
args = {'file': args.file, 'folder': args.dir,
'ofp': args.ofp, 'odp': args.odp}
return args
class Api:
def __init__(self):
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
def call(self, file, folder, ofp, odp):
"""Method saves the predicted image by taking different parameters."""
if file != None and folder != None:
print('"folder" flag and "file" flag cant be used together')
return
model = self._load_model()
save_path = None
# For a single file
if file != None:
image = self._get_file(file)
output = self._get_model_output(image, model)
name, extension = file.split('.')
save_path = name+'_predicted'+'.'+extension
if ofp:
save_path = os.path.join(ofp,save_path)
self._save_image(output, save_path)
print(f'Output Image Saved At {save_path}')
elif folder != None:
image_list = os.listdir(folder)
for file in image_list:
file_name = os.path.join(folder, file)
image = self._get_file(file_name)
output = self._get_model_output(image, model)
name, extension = file.split('.')
save_path = name+'_predicted'+'.'+extension
save_path = os.path.join(
odp, save_path) if odp else os.path.join(folder, save_path)
self._save_image(output, save_path)
print(f'Output Image Saved At {save_path}')
def _load_model(self):
"""Load the saved model and return it."""
filter_list = [16, 32, 64, 128, 256]
model = DynamicUNet(filter_list).to(self.device)
classifier = BrainTumorClassifier(model, self.device)
model_path = os.path.join(
'saved_models', 'UNet-[16, 32, 64, 128, 256].pt')
classifier.restore_model(model_path)
print(
f'Saved model at location "{model_path}" loaded on {self.device}')
return model
def _get_model_output(self, image, model):
"""Returns the saved model output"""
image = image.view((-1, 1, 512, 512)).to(self.device)
output = model(image).detach().cpu()
output = (output > 0.5)
output = output.numpy()
output = np.resize((output * 255), (512, 512))
return output
def _save_image(self, image, path):
"""Save the image to storage specified by path"""
image = Image.fromarray(np.uint8(image), 'L')
image.save(path)
def _get_file(self, file_name):
"""Load the image by taking file name as input"""
default_transformation = transforms.Compose([
transforms.Grayscale(),
transforms.Resize((512, 512))
])
image = default_transformation(Image.open(file_name))
return TF.to_tensor(image)
if __name__ == "__main__":
args = get_arguments()
api = Api()
api.call(**args)