| 307 | |
| 308 | |
| 309 | def get_layers(model_dir, model_name, layers): |
| 310 | model_file = "%s/%s.pb" % (model_dir, model_name) |
| 311 | output_dir = "%s/output_models/" % model_dir |
| 312 | if os.path.exists(output_dir): |
| 313 | sh.rm('-rf', output_dir) |
| 314 | os.makedirs(output_dir) |
| 315 | |
| 316 | convert(model_file, output_dir, layers) |
| 317 | |
| 318 | output_configs_path = output_dir + "outputs.yml" |
| 319 | with open(output_configs_path) as f: |
| 320 | output_configs = yaml.load(f) |
| 321 | output_configs = output_configs['subgraphs'] |
| 322 | |
| 323 | return output_configs |
| 324 | |
| 325 | |
| 326 | def parse_args(): |