-
Notifications
You must be signed in to change notification settings - Fork 0
/
predictor.py
92 lines (74 loc) · 4.8 KB
/
predictor.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
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import os
from parameters import *
from networks.nets import *
from dataloader.dataset import MetaDataset
from dataloader import dict_transforms
import functions.utils as util
if __name__ == "__main__":
# ------------------------------------------------------------------------------------------------------------------
# ========================================== dir & param ==========================================
branch_num = 'branch_7'
model_name = r'ResNet101-DeepLabV3_epoch_350.pth'
data_dir = r'/home/user/Desktop/test_dataset'
weight_dir = os.path.join(r'./save', branch_num, model_name)
dst_dir = os.path.join(data_dir, 'predicton/', branch_num)
store_num = 10
the_name = os.path.splitext(os.path.basename(weight_dir))[0]
assert test_params['test_batch'] >= store_num, 'batch size must be bigger than the number of storing image.'
# =================================================================================================
# ------------------------------------------------------------------------------------------------------------------
# =========================================== transform ===========================================
transform_set1 = transforms.Compose([dict_transforms.Resize(params['resized']),
dict_transforms.Normalize(mean=params['mean'], std=params['std']),
dict_transforms.ToTensor()])
transform_set2 = transforms.Compose([dict_transforms.AlphaKill(),
dict_transforms.Resize(params['resized']),
dict_transforms.Normalize(mean=params['mean'], std=params['std']),
dict_transforms.ToTensor()])
# =================================================================================================
# ------------------------------------------------------------------------------------------------------------------
# ============================================ Dataset ============================================
pred_data1 = MetaDataset(data_dir=data_dir, transform=transform_set1, dataname_extension='*.jpg')
pred_data2 = MetaDataset(data_dir=data_dir, transform=transform_set2, dataname_extension='*.png')
pred_data = pred_data1 + pred_data2
data_loader = DataLoader(dataset=pred_data, batch_size=test_params['test_batch'],
shuffle=False, num_workers=user_setting['test_processes'])
print(f'{len(pred_data)} data detected.')
# =================================================================================================
# ------------------------------------------------------------------------------------------------------------------
# =========================================== Model Load ==========================================
netend = NetEnd(1)
model = ResNet101_DeeplabV3(end_module=netend, pretrain=permission['pretrain'])
model.load_state_dict(torch.load(weight_dir), strict=False)
# =================================================================================================
# ------------------------------------------------------------------------------------------------------------------
# ========================================== GPU setting ==========================================
environment = {}
if torch.cuda.is_available():
device = torch.device('cuda')
print(f'GPU {torch.cuda.get_device_name()} available.')
model.cuda()
environment['gpu'] = True
else:
device = torch.device('cpu')
print(f'GPU unable.')
environment['gpu'] = False
# =================================================================================================
# ------------------------------------------------------------------------------------------------------------------
os.makedirs(dst_dir, exist_ok=True)
print(f'save directory : {dst_dir}')
# ============================================ run ================================================
model.eval()
for i, data in enumerate(data_loader):
image = data[tag_image]
name = data[tag_name]
if environment['gpu']:
image = image.cuda()
with torch.no_grad():
output = model.forward(image)
util.imgstore(output*255.0, nums=store_num, save_dir=dst_dir, epoch=the_name, filename=name)
print(f'{(i / len(data_loader)) * 100:.2f} % done.')
# =================================================================================================
# ------------------------------------------------------------------------------------------------------------------