-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference_mask.py
92 lines (70 loc) · 2.99 KB
/
inference_mask.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
import argparse
import multiprocessing
import os
from importlib import import_module
import pandas as pd
import torch
from torch.utils.data import DataLoader
from dataset_mask import TestDataset
def load_model(saved_model, num_classes, device):
model_cls = getattr(import_module("model_mask"), args.model)
model = model_cls(
num_classes=num_classes
)
# tarpath = os.path.join(saved_model, 'best.tar.gz')
# tar = tarfile.open(tarpath, 'r:gz')
# tar.extractall(path=saved_model)
model_path = os.path.join(saved_model, 'best.pth')
model.load_state_dict(torch.load(model_path, map_location=device))
return model
@torch.no_grad()
def inference(data_dir, model_dir, output_dir, args):
"""
"""
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
num_classes = 3
model = load_model(model_dir, num_classes, device).to(device)
model.eval()
img_root = os.path.join(data_dir, 'images')
info_path = os.path.join(data_dir, 'info.csv')
info = pd.read_csv(info_path)
img_paths = [os.path.join(img_root, img_id) for img_id in info.ImageID]
dataset = TestDataset(img_paths, args.resize)
loader = DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=multiprocessing.cpu_count() // 2,
shuffle=False,
pin_memory=use_cuda,
drop_last=False,
)
print("Calculating inference results..")
preds = []
with torch.no_grad():
for idx, images in enumerate(loader):
images = images.to(device)
pred = model(images)
pred = pred.argmax(dim=-1)
preds.extend(pred.cpu().numpy())
info['ans'] = preds
save_path = os.path.join(output_dir, f'{args.output_name}.csv')
info.to_csv(save_path, index=False)
print(f"Inference Done! Inference result saved at {save_path}")
if __name__ == '__main__':
parser = argparse.ArgumentParser()
# Data and model checkpoints directories
parser.add_argument('--batch_size', type=int, default=64, help='input batch size for validing (default: 1000)')
parser.add_argument('--resize', type=tuple, default=(256, 192), help='resize size for image when you trained (default: (96, 128))')
parser.add_argument('--model', type=str, default='ModelMask', help='model type (default: BaseModel)')
# Container environment
parser.add_argument('--data_dir', type=str, default=os.environ.get('SM_CHANNEL_EVAL', '/opt/ml/input/data/eval'))
parser.add_argument('--model_dir', type=str, default=os.environ.get('SM_CHANNEL_MODEL', './model/exp_mask'))
parser.add_argument('--output_dir', type=str, default=os.environ.get('SM_OUTPUT_DATA_DIR', './output_mask'))
parser.add_argument('--output_name', type=str, default='output_mask')
args = parser.parse_args()
data_dir = args.data_dir
model_dir = args.model_dir
output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)
inference(data_dir, model_dir, output_dir, args)