RangeRet/infer.py at main · SiMoM0/RangeRet · GitHub

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
#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import argparse
import subprocess
import datetime
import yaml
from shutil import copyfile
import os
import shutil

from modules.user import User

if __name__ == '__main__':
	splits = ('train', 'valid', 'test')
	parser = argparse.ArgumentParser("./infer.py")
	parser.add_argument(
		'--dataset', '-d',
		type=str,
		required=True,
		help='Dataset to train with. No Default',
	)
	parser.add_argument(
		'--config',
		type=str,
		required=False,
		default='config/RangeRet-semantickitti.yaml',
		help='Architecture yaml cfg file. See /config/ for sample. No default!',
	)
	parser.add_argument(
		'--data',
		type=str,
		required=False,
		default='config/labels/semantic-kitti.yaml',
		help='Classification yaml cfg file. See /config/labels for sample. No default!',
	)
	parser.add_argument(
		'--log', '-l',
		type=str,
		default=os.getcwd() + '/log/preds/',
		help='Directory to put the predictions. Default: log/preds'
	)
	parser.add_argument(
		'--model', '-m',
		type=str,
		required=True,
		default=None,
		help='Directory to get the trained model.'
	)
	parser.add_argument(
        '--split',
        type=str,
        required=True,
        default='valid',
        help='Split to evaluate on. One of ' +
             str(splits) + '. Defaults to %(default)s',
    )
	parser.add_argument(
        '--save',
        action='store_true',
		default=False,
        help='Save predictions in the log directory. Default: False',
    )
	parser.add_argument(
		'--fp16',
		action='store_true',
		default=False,
		help='Use FP16 precision for inference. Default: False',
	)
	FLAGS, unparsed = parser.parse_known_args()

	# print summary of what we will do
	print("----------")
	print("INTERFACE:")
	print("dataset", FLAGS.dataset)
	print("config", FLAGS.config)
	print("data", FLAGS.data)
	print("log", FLAGS.log)
	print("model", FLAGS.model)
	print("split", FLAGS.split)
	print("save", FLAGS.save)
	print("fp16", FLAGS.fp16)
	print("----------\n")
	# print("Commit hash (training version): ", str(
	# 	subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).strip()))
	print("----------\n")

	# open arch config file
	try:
		print("Opening arch config file from %s" % FLAGS.config)
		ARCH = yaml.safe_load(open(FLAGS.config,'r'))
	except Exception as e:
		print(e)
		print("Error opening arch yaml file.")
		quit()

	# open data config file
	try:
		print("Opening data config file from %s" % FLAGS.data)
		DATA = yaml.safe_load(open(FLAGS.data, 'r'))
	except Exception as e:
		print(e)
		print("Error opening data yaml file.")
		quit()

	# create log folder
	try:
		if os.path.isdir(FLAGS.log):
			shutil.rmtree(FLAGS.log)
		os.makedirs(FLAGS.log)
		os.makedirs(os.path.join(FLAGS.log, "sequences"))
		# Create folders based on split
		for seq in DATA["split"][FLAGS.split]:
			seq = '{0:02d}'.format(int(seq))
			os.makedirs(os.path.join(FLAGS.log, "sequences", seq))
			os.makedirs(os.path.join(FLAGS.log, "sequences", seq, "predictions"))
	except Exception as e:
		print(e)
		print("Error creating log directory. Check permissions!")
		raise

	# does model folder exist?
	if os.path.isfile(FLAGS.model):
		print("model exists! Using model from %s" % (FLAGS.model))
	else:
		print("model does not exist! Can't infer...")
		quit()

	# create user and infer dataset
	user = User(ARCH, DATA, FLAGS.dataset, FLAGS.log, FLAGS.model, FLAGS.split, FLAGS.save, FLAGS.fp16)
	user.infer()