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
|
#!/usr/bin/python3
import glob
import hashlib
import json
import os
import readline
import shutil
import subprocess
def init_cli():
def complete(text, state):
return (glob.glob(text+'*')+[None])[state]
readline.set_completer_delims(' \t\n;')
readline.parse_and_bind("tab: complete")
readline.set_completer(complete)
def get_sha256(filepath):
bufsize = 65536
sha256 = hashlib.sha256()
with open(filepath, 'rb') as f:
while True:
data = f.read(bufsize)
if not data:
break
sha256.update(data)
return sha256.hexdigest()
def get_meta():
"""Get metadata from user via CLI.
:return: metadata in JSON structure
:rtype: dict
"""
inference_engine_names = {
1: 'tflite',
2: 'openvino',
3: 'darknet',
4: 'tensorflow',
5: 'pytorch',
6: 'keras',
7: 'mxnet',
8: 'caffe2',
9: 'caffe',
10: 'movidius',
11: 'others'
}
# model, label, config are full paths for copying them into
# target package directory. They will be updated to relative paths
# in the model package.
meta = {}
config_files = {}
meta['name'] = input('Package name: ')
meta['version'] = input('Package version: ')
meta['model'] = os.path.abspath(input('Model filepath: '))
meta['label'] = os.path.abspath(input('Label filepath: '))
while True:
key = input('Config name (press enter directly to stop): ')
if len(key) != 0:
value = os.path.abspath(input('Config filepath: '))
config_files[key] = value
else:
break
meta['config'] = config_files
engine_index = int(input(
(
'Inference engine\n'
'\t 1. TFLite\n'
'\t 2. OpenVINO\n'
'\t 3. Darknet\n'
'\t 4. TensorFlow\n'
'\t 5. PyTorch\n'
'\t 6. Keras\n'
'\t 7. MXNet\n'
'\t 8. Caffe2\n'
'\t 9. Caffe\n'
'\t10. Movidius\n'
'\t11. Others\n'
': '
)
))
meta['inference-engine'] = inference_engine_names[engine_index]
return meta
def create_metafile(meta, package_dirpath):
"""Create DL model package meta file (meta.json)
"""
checksums = {}
for cksum_key in ['model', 'label']:
target_path = os.path.join(package_dirpath, meta[cksum_key])
checksums[meta[cksum_key]] = get_sha256(target_path)
for k, v in meta['config'].items():
target_path = os.path.join(package_dirpath, v)
checksums[v] = get_sha256(target_path)
meta['checksums-sha256'] = checksums
with open(os.path.join(package_dirpath, 'meta.json'), 'w') as f:
json.dump(meta, f, indent=4)
def create_source_package(meta):
# copy model contents to model package directory
package_name = meta['name'] + '-' + meta['version']
package_dirpath = os.path.join('/tmp', package_name)
subprocess.call(
'mkdir -p {pkgdir}'.format(pkgdir=package_dirpath),
shell=True)
shutil.copy2(meta['model'], package_dirpath)
shutil.copy2(meta['label'], package_dirpath)
for k, v in meta['config'].items():
shutil.copy2(v, package_dirpath)
# update full paths to relative paths in the model package directory
meta['model'] = os.path.basename(meta['model'])
meta['label'] = os.path.basename(meta['label'])
for k, v in meta['config'].items():
meta['config'][k] = os.path.basename(v)
# create model description file (meta.json)
create_metafile(meta, package_dirpath)
print('Model source package is at ' + package_dirpath)
def main():
init_cli()
meta = get_meta()
create_source_package(meta)
if __name__ == '__main__':
# TODO: capture ctrl+c and confirm exit or not
main()
|
