-
Notifications
You must be signed in to change notification settings - Fork 1
/
inference.py
122 lines (103 loc) · 4.37 KB
/
inference.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
114
115
116
117
118
119
120
121
122
import argparse
from tqdm import tqdm
import os
import importlib
from pathlib import Path
import pickle
import numpy as np
from collections import defaultdict
import albumentations as albu
import torch
from pneumothorax_dataset import PneumothoraxDataset
from utils.helpers import load_yaml, init_seed, init_logger
def argparser():
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, help='Path to config file path')
return vars(parser.parse_args())
def build_checkpoints_list(cfg):
pipeline_path = Path(cfg['CHECKPOINTS']['PIPELINE_PATH'])
pipeline_name = cfg['CHECKPOINTS']['PIPELINE_NAME']
checkpoints_list = []
if cfg.get('SUBMIT_BEST', False):
best_checkpoint_folder = Path(pipeline_path, cfg['CHECKPOINTS']['BEST_FOLDER'])
usefolds = cfg['USEFOLDS']
for fold_id in usefolds:
filename = f'{pipeline_name}_fold_{fold_id}.pth'
checkpoints_list.append(Path(best_checkpoint_folder, filename))
else:
folds_dict = cfg['SELECTED_CHECKPOINTS']
for folder_name, epoch_list in folds_dict.items():
checkpoint_folder = Path(
pipeline_path,
cfg['CHECKPOINTS']['FULL_FOLDER'],
folder_name
)
for epoch in epoch_list:
checkpoint_path = Path(
checkpoint_folder,
f'{pipeline_name}_{folder_name}_{epoch:03d}.pth'
)
checkpoints_list.append(checkpoint_path)
return checkpoints_list
def inference_image(model, images, device):
images = images.to(device)
preds = model(images)
masks = torch.sigmoid(preds)
masks = masks.squeeze(1).cpu().detach().numpy()
return masks
def inference_model(model, loader, device, use_flip, result_path, split=100):
mask_dict = {}
result_path = str(result_path)
counter = 0
for image_ids, images in tqdm(loader):
masks = inference_image(model, images, device)
if use_flip:
flipped_images = torch.flip(images, dims=(3,))
flipped_masks = inference_image(model, flipped_images, device)
flipped_masks = np.flip(flipped_masks, axis=2)
masks = (masks + flipped_masks) / 2
for name, mask in zip(image_ids, masks):
mask_dict[name] = mask.astype(np.float32)
counter += 1
if counter % split == 0 or counter == len(loader):
prefix = result_path[:result_path.rfind('.')]
path = f'{prefix}_{(counter - 1) // split + 1:02d}.pkl'
with open(path, 'wb') as handle:
pickle.dump(mask_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
mask_dict = {}
return None
if __name__ == '__main__':
args = argparser()
config_path = Path(args['config'].strip('/'))
experiment_folder = config_path.parents[0]
inference_config = load_yaml(config_path)
batch_size = inference_config['BATCH_SIZE']
device = inference_config['DEVICE']
module = importlib.import_module(inference_config['MODEL']['PY'])
model_class = getattr(module, inference_config['MODEL']['CLASS'])
model = model_class(**inference_config['MODEL'].get('ARGS', None)).to(device)
num_workers = inference_config['WORKERS']
transform = albu.load(inference_config['TEST_TRANSFORMS'])
dataset_folder = inference_config['DATA_DIRECTORY']
dataset = PneumothoraxDataset(
data_folder=dataset_folder,
mode='test',
transform=transform
)
dataloader = torch.utils.data.DataLoader(
dataset=dataset, batch_size=batch_size,
num_workers=num_workers, shuffle=False
)
use_flip = inference_config['FLIP']
checkpoints_list = build_checkpoints_list(inference_config)
result_path = Path(experiment_folder, inference_config['RESULT'])
mask_dict = defaultdict(int)
for pred_idx, checkpoint_path in enumerate(checkpoints_list):
print(checkpoint_path)
model.load_state_dict(torch.load(checkpoint_path))
model.eval()
current_mask_dict = inference_model(model, dataloader, device, use_flip, result_path)
# for name, mask in current_mask_dict.items():
# mask_dict[name] = (mask_dict[name] * pred_idx + mask) / (pred_idx + 1)
# with open(result_path, 'wb') as handle:
# pickle.dump(mask_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)