-
Notifications
You must be signed in to change notification settings - Fork 152
/
test_models.py
228 lines (182 loc) · 7.25 KB
/
test_models.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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
#
#
# 0=================================0
# | Kernel Point Convolutions |
# 0=================================0
#
#
# ----------------------------------------------------------------------------------------------------------------------
#
# Callable script to start a training on ModelNet40 dataset
#
# ----------------------------------------------------------------------------------------------------------------------
#
# Hugues THOMAS - 06/03/2020
#
# ----------------------------------------------------------------------------------------------------------------------
#
# Imports and global variables
# \**********************************/
#
# Common libs
import signal
import os
import numpy as np
import sys
import torch
# Dataset
from datasets.ModelNet40 import *
from datasets.S3DIS import *
from datasets.SensatUrban import *
from datasets.SemanticKitti import *
from datasets.Toronto3D import *
from torch.utils.data import DataLoader
from utils.config import Config
from utils.tester import ModelTester
from models.architectures import KPCNN, KPFCNN
# ----------------------------------------------------------------------------------------------------------------------
#
# Main Call
# \***************/
#
def model_choice(chosen_log):
###########################
# Call the test initializer
###########################
# Automatically retrieve the last trained model
if chosen_log in ['last_ModelNet40', 'last_ShapeNetPart', 'last_S3DIS', 'last_sensaturban']:
# Dataset name
test_dataset = '_'.join(chosen_log.split('_')[1:])
# List all training logs
logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')])
# Find the last log of asked dataset
for log in logs[::-1]:
log_config = Config()
log_config.load(log)
if log_config.dataset.startswith(test_dataset):
chosen_log = log
break
if chosen_log in ['last_ModelNet40', 'last_ShapeNetPart', 'last_S3DIS', 'last_SensatUrban']:
raise ValueError('No log of the dataset "' + test_dataset + '" found')
# Check if log exists
if not os.path.exists(chosen_log):
raise ValueError('The given log does not exists: ' + chosen_log)
return chosen_log
# ----------------------------------------------------------------------------------------------------------------------
#
# Main Call
# \***************/
#
if __name__ == '__main__':
###############################
# Choose the model to visualize
###############################
# Here you can choose which model you want to test with the variable test_model. Here are the possible values :
#
# > 'last_XXX': Automatically retrieve the last trained model on dataset XXX
# > '(old_)results/Log_YYYY-MM-DD_HH-MM-SS': Directly provide the path of a trained model
chosen_log = 'results/Log_2024-05-14_21-04-36'
# Choose the index of the checkpoint to load OR None if you want to load the current checkpoint
chkp_idx = -1
# Choose to test on validation or test split
on_val = True
# Deal with 'last_XXXXXX' choices
chosen_log = model_choice(chosen_log)
############################
# Initialize the environment
############################
# Set which gpu is going to be used
GPU_ID = '0'
# Set GPU visible device
os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID
###############
# Previous chkp
###############
# Find all checkpoints in the chosen training folder
chkp_path = os.path.join(chosen_log, 'checkpoints')
chkps = [f for f in os.listdir(chkp_path) if f[:4] == 'chkp']
# Find which snapshot to restore
if chkp_idx is None:
chosen_chkp = 'current_chkp.tar'
else:
chosen_chkp = np.sort(chkps)[chkp_idx]
chosen_chkp = os.path.join(chosen_log, 'checkpoints', chosen_chkp)
# Initialize configuration class
config = Config()
config.load(chosen_log)
##################################
# Change model parameters for test
##################################
# Change parameters for the test here. For example, you can stop augmenting the input data.
#config.augment_noise = 0.0001
#config.augment_symmetries = False
#config.batch_num = 3
#config.in_radius = 4
config.validation_size = 200
config.input_threads = 10
##############
# Prepare Data
##############
print()
print('Data Preparation')
print('****************')
if on_val:
set = 'validation'
else:
set = 'test'
# Initiate dataset
if config.dataset == 'ModelNet40':
test_dataset = ModelNet40Dataset(config, train=False)
test_sampler = ModelNet40Sampler(test_dataset)
collate_fn = ModelNet40Collate
elif config.dataset == 'S3DIS':
test_dataset = S3DISDataset(config, set='validation', use_potentials=True)
test_sampler = S3DISSampler(test_dataset)
collate_fn = S3DISCollate
elif config.dataset == 'SensatUrban':
test_dataset = SensatUrbanDataset(config, set='validation', use_potentials=True)
test_sampler = SensatUrbanSampler(test_dataset)
collate_fn = SensatUrbanCollate
elif config.dataset == 'Toronto3D':
test_dataset = Toronto3DDataset(config, set='test', use_potentials=True)
test_sampler = Toronto3DSampler(test_dataset)
collate_fn = Toronto3DCollate
elif config.dataset == 'SemanticKitti':
test_dataset = SemanticKittiDataset(config, set=set, balance_classes=False)
test_sampler = SemanticKittiSampler(test_dataset)
collate_fn = SemanticKittiCollate
else:
raise ValueError('Unsupported dataset : ' + config.dataset)
# Data loader
test_loader = DataLoader(test_dataset,
batch_size=1,
sampler=test_sampler,
collate_fn=collate_fn,
num_workers=config.input_threads,
pin_memory=True)
# Calibrate samplers
test_sampler.calibration(test_loader, verbose=True)
print('\nModel Preparation')
print('*****************')
# Define network model
t1 = time.time()
if config.dataset_task == 'classification':
net = KPCNN(config)
elif config.dataset_task in ['cloud_segmentation', 'slam_segmentation']:
net = KPFCNN(config, test_dataset.label_values, test_dataset.ignored_labels)
else:
raise ValueError('Unsupported dataset_task for testing: ' + config.dataset_task)
# Define a visualizer class
tester = ModelTester(net, chkp_path=chosen_chkp)
print('Done in {:.1f}s\n'.format(time.time() - t1))
print('\nStart test')
print('**********\n')
# Training
if config.dataset_task == 'classification':
tester.classification_test(net, test_loader, config)
elif config.dataset_task == 'cloud_segmentation':
tester.cloud_segmentation_test(net, test_loader, config)
elif config.dataset_task == 'slam_segmentation':
tester.slam_segmentation_test(net, test_loader, config)
else:
raise ValueError('Unsupported dataset_task for testing: ' + config.dataset_task)